diff --git a/build/libtensor.so b/build/libtensor.so
index 6912a05..e6e11d3 100755
Binary files a/build/libtensor.so and b/build/libtensor.so differ
diff --git a/build/tensor.o b/build/tensor.o
index 6657ec4..a4e4a41 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 532798d..f85759e 100644
--- a/norch/csrc/tensor.cpp
+++ b/norch/csrc/tensor.cpp
@@ -565,14 +565,25 @@ extern "C" {
             exit(-1);
         }
 
+        int ndim = new_ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            shape[i] = new_shape[i];
+        }
+
         // Calculate the total number of elements in the new shape
-        int new_size = 1;
+        int size = 1;
         for (int i = 0; i < new_ndim; i++) {
-            new_size *= new_shape[i];
+            size *= shape[i];
         }
 
         // Check if the total number of elements matches the current tensor's size
-        if (new_size != tensor->size) {
+        if (size != tensor->size) {
             fprintf(stderr, "Cannot reshape tensor. Total number of elements in new shape does not match the current size of the tensor.\n");
             exit(1);
         }
@@ -582,7 +593,7 @@ extern "C" {
             float* result_data;
             cudaMalloc((void **)&result_data, tensor->size * sizeof(float));
             assign_tensor_cuda(tensor, result_data);
-            return create_tensor(result_data, new_shape, new_ndim, device);
+            return create_tensor(result_data, shape, ndim, device);
         } 
         else {
             float* result_data = (float*)malloc(tensor->size * sizeof(float));
@@ -591,7 +602,7 @@ extern "C" {
                 exit(1);
             }
             assign_tensor_cpu(tensor, result_data);
-            return create_tensor(result_data, new_shape, new_ndim, device);
+            return create_tensor(result_data, shape, ndim, device);
         }
     }
 
diff --git a/test.py b/test.py
index 093d5aa..e86e5c2 100644
--- a/test.py
+++ b/test.py
@@ -28,39 +28,19 @@ if __name__ == "__main__":
         [[7.890, 8.901], [9.012, 1.234], [2.345, 3.456]]
     ], requires_grad=True)
 
-    b = norch.Tensor([[
-        [1.234, 2.123, 1.5],
-        [5.678, 6.789, 1.293],
-        [3.635, 4.456, 1.0202],
-        [7.890, 8.901, 1.91],
-    ],[
-        [1.234, 2.123, 1.5],
-        [5.678, 6.789, 1.293],
-        [3.635, 4.456, 1.0202],
-        [7.890, 8.901, 1.91],
-    ],[
-        [1.234, 2.123, 1.5],
-        [5.678, 6.789, 1.293],
-        [3.635, 4.456, 1.0202],
-        [7.890, 8.901, 1.91],
-    ],[
-        [1.234, 2.123, 1.5],
-        [5.678, 6.789, 1.293],
-        [3.635, 4.456, 1.0202],
-        [7.890, 8.901, 1.91],
-    ],[
-        [1.234, 2.123, 1.5],
-        [5.678, 6.789, 1.293],
-        [3.635, 4.456, 1.0202],
-        [7.890, 8.901, 5.91],
-    ]])
-
     b = norch.Tensor([
-        [1.234, 2.123, 1.5]])
+        [1.234, 2.123, 1.5],
+        [5.678, 6.789, 1.293],
+        [3.635, 4.456, 1.0202],
+        [7.890, 8.901, 1.91],
+    ])
+
+    #b = norch.Tensor([
+    #    [1.234, 2.123, 1.5]])
 
     #print(a.shape)
-    b = a.reshape([2,3,5])
-    print(b.shape)
+    c = a.reshape([2,3,5])
+    print(b @ c)
     #c = result.sum()
     #c.backward()
     #print(a.grad)