fix sum cuda memory corruption

This commit is contained in:
lucasdelimanogueira 2024-05-23 00:03:26 -03:00
parent d093015d15
commit b8ad9ef282

View file

@ -210,7 +210,7 @@ extern "C" {
if (strcmp(tensor->device, "cuda") == 0) {
float* result_data;
cudaMalloc((void**)&result_data, size * sizeof(float));
cudaMalloc((void**)&result_data, tensor->size * sizeof(float));
sum_tensor_cuda(tensor, result_data, axis);
if (keepdim) {