From b8ad9ef282ebd5e415008d85fb96e66770a6da23 Mon Sep 17 00:00:00 2001
From: lucasdelimanogueira <lucasln96@hotmail.com>
Date: Thu, 23 May 2024 00:03:26 -0300
Subject: [PATCH] fix sum cuda memory corruption

---
 norch/csrc/tensor.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp
index 9b08196..8e5ffa6 100644
--- a/norch/csrc/tensor.cpp
+++ b/norch/csrc/tensor.cpp
@@ -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) {