diff --git a/build/cuda.cu.o b/build/cuda.cu.o
index 4ed5a3e..1c85a78 100644
Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ
diff --git a/build/tensor.o b/build/tensor.o
index 47e5749..23f5202 100644
Binary files a/build/tensor.o and b/build/tensor.o differ
diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc
index 718f29e..bbd6f48 100644
Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ
diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu
index fdf2cb5..d38187b 100644
--- a/norch/csrc/cuda.cu
+++ b/norch/csrc/cuda.cu
@@ -141,18 +141,20 @@ __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size
     }
 }
 
-__global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int target_axis, int inner_size, int outer_size) {
+__global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size) {
     int tid = blockIdx.x * blockDim.x + threadIdx.x;
 
-    if (tid < outer_size) {
-        int outer_index = tid / inner_size;
-        int inner_index = tid % inner_size;
+    if (tid < result_size) {
+        for (int i = 0; i < shape[axis]; i++) {
+            int index = 0;
+            int remainder = tid;
+            for (int k = ndim - 2; k >= 0; k--) {
+                index += (remainder % shape[k < axis ? k : k + 1]) * strides[k < axis ? k : k + 1];
+                remainder /= shape[k < axis ? k : k + 1];
+            }
+            index += i * axis_stride;
 
-        int offset = outer_index * strides[0] + inner_index;
-
-        for (int i = 0; i < strides[target_axis]; ++i) {
-            int index = offset + i * strides[target_axis + 1];
-            //atomicAdd(&result_data[outer_index * inner_size + inner_index], data[index]);
+            atomicAdd(&result_data[tid], data[index]);
         }
     }
 }
@@ -184,20 +186,30 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
         cudaDeviceSynchronize();
         
     } else {
+        int axis_stride = tensor->strides[axis];
 
-        int target_axis_stride = tensor->strides[axis];
-        int inner_size = tensor->strides[axis + 1];
-        int outer_size = tensor->size / target_axis_stride;
+        // Calculate the size of the resulting tensor
+        int result_size = 1;
+        for (int i = 0; i < tensor->ndim; i++) {
+            if (i != axis) {
+                result_size *= tensor->shape[i];
+            }
+        }
 
+        // Allocate memory for strides and shape on the device
         int* d_strides;
-        cudaMalloc(&d_strides, (tensor->ndim + 1) * sizeof(int));
-        cudaMemcpy(d_strides, tensor->strides, (tensor->ndim + 1) * sizeof(int), cudaMemcpyHostToDevice);
+        int* d_shape;
+        cudaMalloc(&d_strides, tensor->ndim * sizeof(int));
+        cudaMalloc(&d_shape, tensor->ndim * sizeof(int));
+        cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
+        cudaMemcpy(d_shape, tensor->shape, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
 
-        cudaMemset(result_data, 0, outer_size * sizeof(float));
+        // Initialize result_data to 0
+        cudaMemset(result_data, 0, result_size * sizeof(float));
 
-        int num_threads = outer_size;
+        int num_threads = result_size;
         int num_blocks = (num_threads + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
-        sum_tensor_cuda_kernel_axis<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, d_strides, axis, inner_size, outer_size);
+        sum_tensor_cuda_kernel_axis<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, d_strides, d_shape, axis, tensor->ndim, axis_stride, tensor->size, result_size);
 
         cudaError_t error = cudaGetLastError();
         if (error != cudaSuccess) {
@@ -209,6 +221,7 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
 
         // Free allocated memory
         cudaFree(d_strides);
+        cudaFree(d_shape);
     }
 }
 
diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h
index 822d661..2ba86e8 100644
--- a/norch/csrc/cuda.h
+++ b/norch/csrc/cuda.h
@@ -14,7 +14,7 @@
     __host__ void sub_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
 
     __global__ void sum_tensor_cuda_kernel(float* data, float* result_data);
-    __global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int target_axis, int inner_size, int outer_size);
+    __global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size);
     __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis);
 
     __global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp
index 129a747..41403f0 100644
--- a/norch/csrc/tensor.cpp
+++ b/norch/csrc/tensor.cpp
@@ -202,9 +202,9 @@ extern "C" {
             ndim = tensor->ndim - 1;
         }
 
-        int size = 1;
+        int axis_size = 1;
         for (int i = 0; i < ndim; i++) {
-            size *= shape[i];
+            axis_size *= shape[i];
         }
   
         if (strcmp(tensor->device, "cuda") == 0) {
@@ -213,11 +213,7 @@ extern "C" {
             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));
+                cudaMalloc((void**)&result_data, axis_size * sizeof(float));
             }
             sum_tensor_cuda(tensor, result_data, axis);
             
@@ -241,13 +237,13 @@ extern "C" {
             return create_tensor(result_data, shape, ndim, device);
         } 
         else {
-            float* result_data = (float*)calloc(size, sizeof(float));
+            float* result_data = (float*)calloc(axis_size, sizeof(float));
             if (result_data == NULL) {
                 fprintf(stderr, "Memory allocation failed\n");
                 exit(1);
             }
 
-            sum_tensor_cpu(tensor, result_data, size, shape, axis);
+            sum_tensor_cpu(tensor, result_data, axis_size, shape, axis);
 
             if (keepdim) {
                 if (axis == -1){
diff --git a/norch/libtensor.so b/norch/libtensor.so
index eb4a5ac..e70bc77 100755
Binary files a/norch/libtensor.so and b/norch/libtensor.so differ
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..94f5195
--- /dev/null
+++ b/test.py
@@ -0,0 +1,7 @@
+import norch
+from norch.utils import utils_unittests as utils
+
+device = "cpu"
+
+norch_tensor = norch.Tensor([[[[1, 2], [3, -4]], [[5, 6], [7, 8]]], [[[1, 2], [3, -4]], [[5, 6], [7, 8]]]]).to(device)
+norch_result = norch_tensor.sum(axis=0)