diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu
index 9304207..bf413e3 100644
--- a/norch/csrc/cuda.cu
+++ b/norch/csrc/cuda.cu
@@ -4,8 +4,8 @@
 #include <math.h>
 
 #define THREADS_PER_BLOCK 128
-#define THREADS_PER_BLOCK_SUM 1024
 #define TILE_SIZE 32
+#define SHMEM_SIZE THREADS_PER_BLOCK_SUM * sizeof(float)
 
 __host__ void cpu_to_cuda(Tensor* tensor) {
     
@@ -61,7 +61,7 @@ __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_da
 
 
 __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size) {
-    __shared__ float partial_sum[THREADS_PER_BLOCK_SUM];
+    __shared__ float partial_sum[SHMEM_SIZE];
 
     int tid = threadIdx.x;
     int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -84,46 +84,18 @@ __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size
     }
 }
 
-
-__global__ void aux_final_sum_kernel(float* result_data, int size) {
-    __shared__ float partial_sum[THREADS_PER_BLOCK_SUM];
-
-    int tid = threadIdx.x;
-    int i = blockIdx.x * blockDim.x + threadIdx.x;
-
-    partial_sum[tid] = (i < size) ? result_data[i] : 0;
-
-    __syncthreads();
-
-    // Perform final reduction
-    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
-        if (tid < s) {
-            partial_sum[tid] += partial_sum[tid + s];
-        }
-        __syncthreads();
-    }
-
-    // Write final result to global memory
-    if (tid == 0 && blockIdx.x == 0) {
-        result_data[0] = partial_sum[0];
-    }
-}
-
-
-
-
 __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data) {
     cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);
 
-    int num_blocks = (tensor->size + THREADS_PER_BLOCK_SUM - 1) / THREADS_PER_BLOCK_SUM;
+    int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
 
     // First-level reduction
-    sum_tensor_cuda_kernel<<<num_blocks, THREADS_PER_BLOCK_SUM>>>(tensor->data, result_data, tensor->size);
+    sum_tensor_cuda_kernel<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
 
     // If necessary, perform multiple levels of reduction
     while (num_blocks > 1) {
-        int num_blocks_next = (num_blocks + THREADS_PER_BLOCK_SUM - 1) / THREADS_PER_BLOCK_SUM;
-        aux_final_sum_kernel<<<num_blocks_next, THREADS_PER_BLOCK_SUM>>>(result_data, num_blocks);
+        int num_blocks_next = (num_blocks + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+        sum_tensor_cuda_kernel<<<num_blocks_next, THREADS_PER_BLOCK>>>(result_data, result_data, num_blocks);
         num_blocks = num_blocks_next;
     }
 
@@ -137,6 +109,8 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data) {
 }
 
 
+
+
 __global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
    
     int i = blockIdx.x * blockDim.x + threadIdx.x;