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

fix sum cuda axis atomicAdd

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-23 00:30:48 -03:00
parent b8ad9ef282
commit 192a50f773
6 changed files with 34 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

10
norch/csrc/tensor.cpp
View file

@ -210,7 +210,15 @@ extern "C" {
if (strcmp(tensor->device, "cuda") == 0) {
float* result_data;
cudaMalloc((void**)&result_data, tensor->size * sizeof(float));
if (axis == -1) {
cudaMalloc((void**)&result_data, tensor->size * sizeof(float));
} else {
int target_axis_stride = tensor->strides[axis];
int outer_size = tensor->size / target_axis_stride;
cudaMalloc((void**)&result_data, (outer_size) * sizeof(float));
}
sum_tensor_cuda(tensor, result_data, axis);
if (keepdim) {