aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix sum tensor reduce cuda

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-01 14:44:09 -03:00
parent 592dacbde9
commit 2fd5b75ff0
4 changed files with 24 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

6
norch/csrc/tensor.cpp
View file

@ -171,9 +171,9 @@ extern "C" {
if (strcmp(tensor->device, "cuda") == 0) {
float* result_data;
int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
cudaMalloc((void**)&result_data, number_of_blocks * sizeof(float));
sum_tensor_cuda(tensor, result_data, number_of_blocks);
cudaMalloc((void**)&result_data, tensor->size * sizeof(float));
result_data[0] = 0;
sum_tensor_cuda(tensor, result_data);
return create_tensor(result_data, shape, ndim, device);
}
else {