fix sum cuda axis

norch/csrc/cuda.cu

@@ -172,9 +172,10 @@ __global__ void sum_tensor_axis_cuda_kernel(float* data, float* result_data, int
 
 
 __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
-    if (axis == -1) {
-        cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);
+    cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);
 
+    if (axis == -1) {
+        
         int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
 
         // First-level reduction
@@ -198,17 +199,10 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
     } else {
         int axis_size = tensor->shape[axis];
 
-        // Allocate memory for temporary storage on the device
-        float* temp_data;
-        cudaMalloc(&temp_data, tensor->size * sizeof(float));
-
-        // Copy tensor data to device
-        cudaMemcpy(temp_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);
-
         int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
 
         // First-level reduction
-        sum_tensor_axis_cuda_kernel<<<num_blocks, THREADS_PER_BLOCK>>>(temp_data, result_data, tensor->size, axis_size);
+        sum_tensor_axis_cuda_kernel<<<num_blocks, THREADS_PER_BLOCK>>>(result_data, result_data, tensor->size, axis_size);
 
         // If necessary, perform multiple levels of reduction
         while (num_blocks > 1) {
@@ -217,9 +211,6 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
             num_blocks = num_blocks_next;
         }
 
-        // Free allocated memory on the device
-        cudaFree(temp_data);
-
         cudaError_t error = cudaGetLastError();
         if (error != cudaSuccess) {
             printf("CUDA error: %s\n", cudaGetErrorString(error));

norch/csrc/tensor.cpp

@@ -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) {