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

fix sum cuda axis

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-22 21:07:09 -03:00
parent 2c7571c371
commit d093015d15
5 changed files with 6 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

5
norch/csrc/tensor.cpp
View file

@ -182,7 +182,7 @@ extern "C" {
}
int ndim;
int* shape;
if (axis > tensor->ndim - 1) {
fprintf(stderr, "Error: axis argument %d must be smaller than tensor dimension %d", axis, tensor->ndim);
}
@ -208,10 +208,9 @@ extern "C" {
}
if (strcmp(tensor->device, "cuda") == 0) {
float* result_data;
cudaMalloc((void**)&result_data, size * sizeof(float));
cudaMemset(result_data, 0, size * sizeof(float));
sum_tensor_cuda(tensor, result_data, axis);
if (keepdim) {