diff --git a/build/cuda.cu.o b/build/cuda.cu.o
index f088d3e..99d9f0f 100644
Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ
diff --git a/build/libtensor.so b/build/libtensor.so
index 7df5103..c24842c 100755
Binary files a/build/libtensor.so and b/build/libtensor.so differ
diff --git a/build/tensor.o b/build/tensor.o
index 2fc99ff..71aef7c 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 690fbbe..5b296b4 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 b4e1571..df810eb 100644
--- a/norch/csrc/cuda.cu
+++ b/norch/csrc/cuda.cu
@@ -104,8 +104,8 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int number_of_
     int levelSize = number_of_blocks;
     while (remaining > 1) {
         int threads = (remaining + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
-        sum_block<<<1, threads, threads * sizeof(float)>>>(result_data, remaining);
-        remaining = (remaining + threadsPerBlock - 1) / threadsPerBlock;
+        aux_sum_block_kernel<<<1, threads, threads * sizeof(float)>>>(result_data, remaining);
+        remaining = (remaining + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
         levelSize = remaining;
     }
 
diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h
index bb06a67..606fb70 100644
--- a/norch/csrc/cuda.h
+++ b/norch/csrc/cuda.h
@@ -7,8 +7,8 @@
     __global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
     __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);
     
+    __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int number_of_blocks);
     __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size);
-    __host__ void sum_tensor_cuda(Tensor* tensor1, float* result_data);
 
     __global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
     __host__ void sub_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);
@@ -36,4 +36,5 @@
 
 
 
+
 #endif /* CUDA_KERNEL_H_ */
diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp
index 20d9f85..4f7d765 100644
--- a/norch/csrc/tensor.cpp
+++ b/norch/csrc/tensor.cpp
@@ -171,9 +171,9 @@ extern "C" {
         if (strcmp(tensor->device, "cuda") == 0) {
 
             float* result_data;
-            int num_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
-            cudaMalloc((void **)&num_of_blocks, 1 * sizeof(float));
-            sum_tensor_cuda(tensor, result_data, int num_of_blocks);
+            int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+            cudaMalloc((void**)&result_data, number_of_blocks * sizeof(float));
+            sum_tensor_cuda(tensor, result_data, number_of_blocks);
             return create_tensor(result_data, shape, ndim, device);
         } 
         else {
diff --git a/test.py b/test.py
index fb6378b..56813ba 100644
--- a/test.py
+++ b/test.py
@@ -30,11 +30,12 @@ if __name__ == "__main__":
     a = norch.Tensor([[1, 2], [1, 2], [1, 2]], requires_grad=True)#.to("cuda")
     b = norch.Tensor([[1, 400, 3], [1, 2, 3]], requires_grad=True)
 
-    c = (b @ a) * 5
-    d = c.sum()
-    d.backward()
+    print(a.sum())
+    #c = (b @ a) * 5
+    #d = c.sum()
+    #d.backward()
 
-    print(a.grad)
+    #print(a.grad)
     
     """#print(a)
     N = 1000