diff --git a/build/cpu.o b/build/cpu.o
index adc39e5..6372b93 100644
Binary files a/build/cpu.o and b/build/cpu.o differ
diff --git a/build/tensor.o b/build/tensor.o
index 0071d08..8931046 100644
Binary files a/build/tensor.o and b/build/tensor.o differ
diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp
index ffeadbc..c308d58 100644
--- a/norch/csrc/cpu.cpp
+++ b/norch/csrc/cpu.cpp
@@ -205,7 +205,7 @@ void log_tensor_cpu(Tensor* tensor, float* result_data) {
     }
 }
 
-void sum_tensor_cpu(Tensor* tensor, float* result_data, int size, int* result_shape, int axis, bool keepdim) {
+void sum_tensor_cpu(Tensor* tensor, float* result_data, int size, int* result_shape, int axis) {
     if (axis == -1) {
         // Sum over all elements
         float sum = 0.0;
@@ -232,9 +232,6 @@ void sum_tensor_cpu(Tensor* tensor, float* result_data, int size, int* result_sh
                 result_data[j] += tensor->data[index + i * axis_stride];
             }
         }
-        for (int j = 0; j < size; j++) {
-            printf("%f", result_data[j]);
-        }
     }
 }
 
diff --git a/norch/csrc/cpu.h b/norch/csrc/cpu.h
index 5fe358c..28404e5 100644
--- a/norch/csrc/cpu.h
+++ b/norch/csrc/cpu.h
@@ -5,7 +5,7 @@
 
 void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
-void sum_tensor_cpu(Tensor* tensor, float* result_data, int size, int* shape, int axis, bool keepdim);
+void sum_tensor_cpu(Tensor* tensor, float* result_data, int size, int* shape, int axis);
 void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
 void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp
index c945689..b717145 100644
--- a/norch/csrc/tensor.cpp
+++ b/norch/csrc/tensor.cpp
@@ -192,24 +192,14 @@ extern "C" {
             shape[0] = 1;
             ndim = 1;
         } else {
-
-            if (keepdim) {
-                shape = (int*) malloc((tensor->ndim) * sizeof(int));
-                for (int i = 0; i < tensor->ndim; i++) {
-                    shape[i] = tensor->shape[i];
+            shape = (int*) malloc((tensor->ndim - 1) * sizeof(int));
+            for (int i = 0, j = 0; i < tensor->ndim; ++i) {
+                if (i != axis) {
+                    shape[j++] = tensor->shape[i];
                 }
-                shape[axis] = 1;
-                ndim = tensor->ndim;
-                
-            } else {
-                shape = (int*) malloc((tensor->ndim - 1) * sizeof(int));
-                for (int i = 0, j = 0; i < tensor->ndim; ++i) {
-                    if (i != axis) {
-                        shape[j++] = tensor->shape[i];
-                    }
-                }
-                ndim = tensor->ndim - 1;
             }
+            ndim = tensor->ndim - 1;
+            
         }
 
         int size = 1;
@@ -230,7 +220,18 @@ extern "C" {
                 fprintf(stderr, "Memory allocation failed\n");
                 exit(1);
             }
-            sum_tensor_cpu(tensor, result_data, size, shape, axis, keepdim);
+            sum_tensor_cpu(tensor, result_data, size, shape, axis);
+
+            if (keepdim) {
+                shape = (int*) malloc((tensor->ndim) * sizeof(int));
+                for (int i = 0; i < tensor->ndim; i++) {
+                    shape[i] = tensor->shape[i];
+                }
+                shape[axis] = 1;
+                ndim = tensor->ndim;
+                
+            }
+
             return create_tensor(result_data, shape, ndim, device);
         }     
     }
diff --git a/norch/libtensor.so b/norch/libtensor.so
index 09af449..9212825 100755
Binary files a/norch/libtensor.so and b/norch/libtensor.so differ