diff --git a/build/tensor.o b/build/tensor.o
index d12bc52..85b704a 100644
Binary files a/build/tensor.o and b/build/tensor.o differ
diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp
index a340c73..63568cc 100644
--- a/norch/csrc/tensor.cpp
+++ b/norch/csrc/tensor.cpp
@@ -183,11 +183,8 @@ extern "C" {
         int ndim;
         int* shape;
         if (axis == -1) {
+            
             shape = (int*) malloc(sizeof(int));
-            if (shape == NULL) {
-                fprintf(stderr, "Memory allocation failed\n");
-                exit(1);
-            }
             shape[0] = 1;
             ndim = 1;
         } else {
@@ -198,7 +195,6 @@ extern "C" {
                 }
             }
             ndim = tensor->ndim - 1;
-            
         }
 
         int size = 1;
@@ -220,25 +216,27 @@ extern "C" {
                 fprintf(stderr, "Memory allocation failed\n");
                 exit(1);
             }
+
             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++) {
-                    if (axis == -1) {
+                if (axis == -1){
+                    ndim = tensor->ndim;
+                    shape = (int*) malloc((tensor->ndim) * sizeof(int));
+                    for (int i = 0; i < tensor->ndim; i++) {
                         shape[i] = 1;
-                    } else {
+                    }
+                } else {
+                    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;
                 }
-                shape[axis] = 1;
-                ndim = tensor->ndim;
-                Tensor* new_tensor = create_tensor(result_data, shape, ndim, device);
-                return reshape_tensor(new_tensor, shape, ndim);
+                
             }
-
             return create_tensor(result_data, shape, ndim, device);
-            
         }     
     }
 
@@ -252,12 +250,8 @@ extern "C" {
         }
         int ndim;
         int* shape;
-        if (axis == -1) {
+        if (axis == -1) {     
             shape = (int*) malloc(sizeof(int));
-            if (shape == NULL) {
-                fprintf(stderr, "Memory allocation failed\n");
-                exit(1);
-            }
             shape[0] = 1;
             ndim = 1;
         } else {
@@ -268,7 +262,6 @@ extern "C" {
                 }
             }
             ndim = tensor->ndim - 1;
-            
         }
 
         int size = 1;
@@ -293,18 +286,20 @@ extern "C" {
             max_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++) {
-                    if (axis == -1) {
+                if (axis == -1){
+                    ndim = tensor->ndim;
+                    shape = (int*) malloc((tensor->ndim) * sizeof(int));
+                    for (int i = 0; i < tensor->ndim; i++) {
                         shape[i] = 1;
-                    } else {
+                    }
+                } else {
+                    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;
-                Tensor* new_tensor = create_tensor(result_data, shape, ndim, device);
-                return reshape_tensor(new_tensor, shape, ndim);
+                    shape[axis] = 1;
+                    ndim = tensor->ndim;
+                }   
             }
 
             return create_tensor(result_data, shape, ndim, device);
@@ -322,11 +317,8 @@ extern "C" {
         int ndim;
         int* shape;
         if (axis == -1) {
+            
             shape = (int*) malloc(sizeof(int));
-            if (shape == NULL) {
-                fprintf(stderr, "Memory allocation failed\n");
-                exit(1);
-            }
             shape[0] = 1;
             ndim = 1;
         } else {
@@ -337,9 +329,8 @@ extern "C" {
                 }
             }
             ndim = tensor->ndim - 1;
-            
         }
-
+        
         int size = 1;
         for (int i = 0; i < ndim; i++) {
             size *= shape[i];
@@ -362,18 +353,20 @@ extern "C" {
             min_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++) {
-                    if (axis == -1) {
+                if (axis == -1){
+                    ndim = tensor->ndim;
+                    shape = (int*) malloc((tensor->ndim) * sizeof(int));
+                    for (int i = 0; i < tensor->ndim; i++) {
                         shape[i] = 1;
-                    } else {
+                    }
+                } else {
+                    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;
-                Tensor* new_tensor = create_tensor(result_data, shape, ndim, device);
-                return reshape_tensor(new_tensor, shape, ndim);
+                    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 5d6e518..de7ec33 100755
Binary files a/norch/libtensor.so and b/norch/libtensor.so differ
diff --git a/tests/test_operations.py b/tests/test_operations.py
index 8a684e9..dffe0a1 100644
--- a/tests/test_operations.py
+++ b/tests/test_operations.py
@@ -272,6 +272,7 @@ class TestTensorOperations(unittest.TestCase):
         torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch.sum(torch_tensor, dim=1, keepdim=True)
         
+        print(torch_result, torch_expected)
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
     def test_max(self):