diff --git a/build/cpu.o b/build/cpu.o
index b6465ab..8eb61f6 100644
Binary files a/build/cpu.o and b/build/cpu.o differ
diff --git a/build/cuda.cu.o b/build/cuda.cu.o
index 51d5b7b..6936037 100644
Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ
diff --git a/build/libtensor.so b/build/libtensor.so
index 1f13ba8..85e0f90 100755
Binary files a/build/libtensor.so and b/build/libtensor.so differ
diff --git a/build/tensor.o b/build/tensor.o
index 1d75a22..4e7c8e8 100644
Binary files a/build/tensor.o and b/build/tensor.o differ
diff --git a/csrc/cpu.cpp b/csrc/cpu.cpp
index 93623e9..61a167c 100644
--- a/csrc/cpu.cpp
+++ b/csrc/cpu.cpp
@@ -25,6 +25,19 @@ void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_
     }
 }
 
+void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+    
+    for (int i = 0; i < tensor1->shape[0]; i++) {
+        for (int j = 0; j < tensor2->shape[1]; j++) {
+            float sum = 0.0;
+            for (int k = 0; k < tensor1->shape[1]; k++) {
+                sum += tensor1->data[i * tensor1->shape[1] + k] * tensor2->data[k * tensor2->shape[1] + j];
+            }
+            result_data[i * tensor2->shape[1] + j] = sum;
+        }
+    }
+}
+
 void pow_tensor_cpu(Tensor* tensor, float power, float* result_data) {
     
     for (int i = 0; i < tensor->size; i++) {
diff --git a/csrc/cpu.h b/csrc/cpu.h
index 20c0f2b..4370dc8 100644
--- a/csrc/cpu.h
+++ b/csrc/cpu.h
@@ -1,5 +1,5 @@
 void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, 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 pow_tensor_cpu(Tensor* tensor, float power, float* result_data);
\ No newline at end of file
diff --git a/csrc/cuda.cu b/csrc/cuda.cu
index d433677..6a953b2 100644
--- a/csrc/cuda.cu
+++ b/csrc/cuda.cu
@@ -6,6 +6,7 @@
 #define THREADS_PER_BLOCK 128
 
 __host__ void cpu_to_cuda(Tensor* tensor) {
+    
     float* data_tmp;
     cudaMalloc((void **)&data_tmp, tensor->size * sizeof(float));
     cudaMemcpy(data_tmp, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);
@@ -35,6 +36,7 @@ __host__ void cuda_to_cpu(Tensor* tensor) {
 }
 
 __global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
+    
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i < size) {
         result_data[i] = data1[i] + data2[i];
@@ -56,6 +58,7 @@ __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_da
 }
 
 __global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
+   
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i < size) {
         result_data[i] = data1[i] - data2[i];
@@ -77,6 +80,7 @@ __host__ void sub_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_da
 }
 
 __global__ void elementwise_mul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
+    
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i < size) {
         result_data[i] = data1[i] * data2[i];
@@ -97,7 +101,43 @@ __host__ void elementwise_mul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, floa
     cudaDeviceSynchronize();
 }
 
+__global__ void matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int rows1, int cols1, int cols2) {    
+
+    int row = blockIdx.y * blockDim.y + threadIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (row < rows1 && col < cols2) {
+        float sum = 0.0;
+        for (int k = 0; k < cols1; k++) {
+            sum += data1[row * cols1 + k] * data2[k * cols2 + col];
+        }
+        result_data[row * cols2 + col] = sum;
+    }
+
+}
+
+__host__ void matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+    
+    int rows1 = tensor1->shape[0];
+    int cols1 = tensor1->shape[1];
+    int cols2 = tensor2->shape[1];
+
+    dim3 threadsPerBlock(16, 16);
+    dim3 numBlocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y);
+    matmul_tensor_cuda_kernel<<<numBlocks, threadsPerBlock>>>(tensor1->data, tensor2->data, result_data, rows1, cols1, cols2);
+
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
 __global__ void pow_tensor_cuda_kernel(float* data, float power, float* result_data, int size) {
+    
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i < size) {
         result_data[i] = powf(data[i], power);
@@ -118,3 +158,4 @@ __host__ void pow_tensor_cuda(Tensor* tensor, float power, float* result_data) {
     cudaDeviceSynchronize();
 }
 
+
diff --git a/csrc/cuda.h b/csrc/cuda.h
index c1a1404..1e83602 100644
--- a/csrc/cuda.h
+++ b/csrc/cuda.h
@@ -16,4 +16,7 @@
     __global__ void pow_tensor_cuda_kernel(float* data, float power, float* result_data, int size);
     __host__ void pow_tensor_cuda(Tensor* tensor, float power, float* result_data);
 
+    __global__ void matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int rows1, int cols1, int cols2)
+    __host__ void matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);
+
 #endif /* CUDA_KERNEL_H_ */
diff --git a/csrc/tensor.cpp b/csrc/tensor.cpp
index 05baef0..c140cd0 100644
--- a/csrc/tensor.cpp
+++ b/csrc/tensor.cpp
@@ -264,9 +264,13 @@ extern "C" {
             exit(1);
         }
 
-        char* device = tensor1->device;
-
-        // Calculate the shape of the result tensor
+        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 = tensor1->ndim + tensor2->ndim - 2;
         int* shape = (int*)malloc(ndim * sizeof(int));
         if (shape == NULL) {
@@ -291,17 +295,22 @@ extern "C" {
             exit(1);
         }
 
-        for (int i = 0; i < tensor1->shape[0]; i++) {
-            for (int j = 0; j < tensor2->shape[1]; j++) {
-                float sum = 0.0;
-                for (int k = 0; k < tensor1->shape[1]; k++) {
-                    sum += tensor1->data[i * tensor1->shape[1] + k] * tensor2->data[k * tensor2->shape[1] + j];
-                }
-                result_data[i * tensor2->shape[1] + j] = sum;
-            }
-        }
+        if (strcmp(tensor1->device, "cuda") == 0) {
 
-        return create_tensor(result_data, shape, ndim, device);
+            float* result_data;
+            cudaMalloc((void **)&result_data, size * sizeof(float));
+            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);
+            }
+            matmul_tensor_cpu(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
     }
 
     Tensor* pow_tensor(Tensor* tensor, float power) {
diff --git a/test.py b/test.py
index 32946a6..62959be 100644
--- a/test.py
+++ b/test.py
@@ -29,7 +29,7 @@ if __name__ == "__main__":
 
     #d = b-c
 
-    b = a ** 2
+    b = a @ a
     print(b)
     #print(a ** 2)