diff --git a/build/cpu.o b/build/cpu.o
index d681b76..f1d237e 100644
Binary files a/build/cpu.o and b/build/cpu.o differ
diff --git a/build/libtensor.so b/build/libtensor.so
index 4118958..c430d19 100755
Binary files a/build/libtensor.so and b/build/libtensor.so differ
diff --git a/build/tensor.o b/build/tensor.o
index 3dbd4ca..6657ec4 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 f6d3f53..15f473b 100644
Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ
diff --git a/norch/autograd/__pycache__/functions.cpython-38.pyc b/norch/autograd/__pycache__/functions.cpython-38.pyc
index 615f472..033c83c 100644
Binary files a/norch/autograd/__pycache__/functions.cpython-38.pyc and b/norch/autograd/__pycache__/functions.cpython-38.pyc differ
diff --git a/norch/autograd/functions.py b/norch/autograd/functions.py
index 5c7ebed..bd52020 100644
--- a/norch/autograd/functions.py
+++ b/norch/autograd/functions.py
@@ -35,8 +35,6 @@ class MatmulBackward:
         x, y = self.input
         return [gradient @ y.T, x.T @ gradient]
         
-
-    
 class PowBackward:
     def __init__(self, x, power):
         self.input = [x]
diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp
index 0365de5..47dfd23 100644
--- a/norch/csrc/cpu.cpp
+++ b/norch/csrc/cpu.cpp
@@ -46,7 +46,7 @@ void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
 }
 
 
-void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+void broadcasted_batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
 
     int tensor2_offset = tensor2->shape[1] * tensor2->shape[2];
     int result_data_offset = tensor1->shape[0] * tensor2->shape[2];
@@ -65,6 +65,26 @@ void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_d
     }
 } 
 
+void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+
+    int tensor1_offset = tensor1->shape[1] * tensor1->shape[2];
+    int tensor2_offset = tensor2->shape[1] * tensor2->shape[2];
+    int result_data_offset = tensor1->shape[1] * tensor2->shape[2];
+
+    for (int batch = 0; batch < tensor2->shape[0]; batch++) {
+    
+        for (int i = 0; i < tensor1->shape[1]; i++) {
+            for (int j = 0; j < tensor2->shape[2]; j++) {
+                float sum = 0.0;
+                for (int k = 0; k < tensor1->shape[2]; k++) {
+                    sum += tensor1->data[(batch * tensor1_offset) + i * tensor1->shape[2] + k] * tensor2->data[batch*tensor2_offset + (k * tensor2->shape[2] + j)];
+                }
+                result_data[(batch * result_data_offset) + (i * tensor2->shape[2] + j)] = sum;
+            }
+        }
+    }
+} 
+
 void pow_tensor_cpu(Tensor* tensor, float power, float* result_data) {
     
     for (int i = 0; i < tensor->size; i++) {
diff --git a/norch/csrc/cpu.h b/norch/csrc/cpu.h
index b6275b3..ae3eae5 100644
--- a/norch/csrc/cpu.h
+++ b/norch/csrc/cpu.h
@@ -8,6 +8,7 @@ void sum_tensor_cpu(Tensor* tensor1, float* result_data);
 void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
+void broadcasted_batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void pow_tensor_cpu(Tensor* tensor, float power, float* result_data);
 void scalar_mul_tensor_cpu(Tensor* tensor, float scalar, float* result_data);
diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp
index ed4eeb7..532798d 100644
--- a/norch/csrc/tensor.cpp
+++ b/norch/csrc/tensor.cpp
@@ -388,11 +388,11 @@ extern "C" {
         }
     }
 
-    Tensor* batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) {
+    Tensor* broadcasted_batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) {
         //MxN @ BATCHxNxP = BATCHxMxP
         // Check if tensors have compatible shapes for matrix multiplication
         if (tensor1->shape[1] != tensor2->shape[1]) {
-            fprintf(stderr, "Incompatible shapes for matrix multiplication %dx%d and %dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1]);
+            fprintf(stderr, "Incompatible shapes for broadcasted batched matrix multiplication %dx%d and %dx%dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1], tensor2->shape[2]);
             exit(1);
         }
 
@@ -435,7 +435,75 @@ extern "C" {
 
             float* result_data;
             cudaMalloc((void **)&result_data, size * sizeof(float));
-            matmul_tensor_cuda(tensor1, tensor2, result_data);
+            ////broadcasted_batched_matmul_tensor_cuda(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            broadcasted_batched_matmul_tensor_cpu(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+
+    }
+
+    Tensor* batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) {
+        //BATCHxMxN @ BATCHxNxP = BATCHxMxP
+        // Check if tensors have compatible shapes for matrix multiplication
+
+        if (tensor1->shape[0] != tensor2->shape[0]) {
+            fprintf(stderr, "Tensors must have same batch dimension for batch matmul %d and %d\n", tensor1->shape[0], tensor2->shape[0]);
+            exit(1);
+        }
+
+        if (tensor1->shape[2] != tensor2->shape[1]) {
+            fprintf(stderr, "Incompatible shapes for matrix multiplication %dx%d and %dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1]);
+            exit(1);
+        }
+        
+        if (strcmp(tensor1->device, tensor2->device) != 0) {
+            fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
+            exit(1);
+        }
+
+        char* device = (char*)malloc(strlen(tensor1->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor1->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+
+        int ndim = 3;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+        
+        shape[0] = tensor2->shape[0];;
+        shape[1] = tensor1->shape[1];
+        shape[2] = tensor2->shape[2];  
+
+        int size = 1;
+        for (int i = 0; i < ndim; i++) {
+            size *= shape[i];
+        }
+
+        float* result_data = (float*)malloc(size * sizeof(float));
+        if (result_data == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        if (strcmp(tensor1->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, size * sizeof(float));
+            //batched_matmul_tensor_cuda(tensor1, tensor2, result_data);
             return create_tensor(result_data, shape, ndim, device);
         } 
         else {
diff --git a/norch/tensor.py b/norch/tensor.py
index 32a5bd4..fa69bd6 100644
--- a/norch/tensor.py
+++ b/norch/tensor.py
@@ -291,8 +291,22 @@ class Tensor:
         return self  
     
     def __matmul__(self, other):
-        if other.ndim == 3:
-            #batched 3D matmul
+        if self.ndim < 3 and other.ndim == 3:
+            #broadcasted 2D x 3D matmul
+
+            Tensor._C.broadcasted_batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.broadcasted_batched_matmul_tensor.restype = ctypes.POINTER(CTensor)
+            
+            result_tensor_ptr = Tensor._C.broadcasted_batched_matmul_tensor(self.tensor, other.tensor)
+
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = [other.shape[0], self.shape[0], other.shape[2]]
+            result_data.ndim = 3
+            result_data.device = self.device
+
+        elif self.ndim == 3 and other.ndim == 3:
+            #broadcasted 3D x 3D matmul
 
             Tensor._C.batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
             Tensor._C.batched_matmul_tensor.restype = ctypes.POINTER(CTensor)
@@ -301,7 +315,7 @@ class Tensor:
 
             result_data = Tensor()
             result_data.tensor = result_tensor_ptr
-            result_data.shape = [other.shape[0], self.shape[0], other.shape[2]]
+            result_data.shape = [other.shape[0], self.shape[1], other.shape[2]]
             result_data.ndim = 3
             result_data.device = self.device
             
diff --git a/test.py b/test.py
index a65593d..27f3ddf 100644
--- a/test.py
+++ b/test.py
@@ -26,17 +26,40 @@ if __name__ == "__main__":
         [[4.567, 5.678], [6.789, 7.890], [8.901, 9.012]],
         [[1.234, 2.345], [3.456, 4.567], [5.678, 6.789]],
         [[7.890, 8.901], [9.012, 1.234], [2.345, 3.456]]
-    ])
+    ], requires_grad=True)
 
-    b = norch.Tensor([
+    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],
+    ]])
 
     result = b @ a
     print(result)
+    #c = result.sum()
+    #c.backward()
+    #print(a.grad)
     
     #a = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda")
     #b = Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda")