diff --git a/build/cuda.cu.o b/build/cuda.cu.o
index 87633ae..5c6f27d 100644
Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ
diff --git a/norch/autograd/__pycache__/functions.cpython-38.pyc b/norch/autograd/__pycache__/functions.cpython-38.pyc
index 1365c91..d48733e 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/csrc/cuda.cu b/norch/csrc/cuda.cu
index 8922346..460baec 100644
--- a/norch/csrc/cuda.cu
+++ b/norch/csrc/cuda.cu
@@ -17,8 +17,6 @@ __host__ void cpu_to_cuda(Tensor* tensor) {
     const char* device_str = "cuda";
     tensor->device = (char*)malloc(strlen(device_str) + 1);
     strcpy(tensor->device, device_str); 
-
-    printf("Successfully sent tensor to: %s\n", tensor->device);
 }
 
 __host__ void cuda_to_cpu(Tensor* tensor) {
@@ -58,56 +56,56 @@ __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_da
     cudaDeviceSynchronize();
 }
 
-__global__ void add_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* shape1, int* shape2, int* broadcasted_shape, int ndim1, int ndim2, int size) {
+__global__ void add_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size) {
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i >= size) return;
-
-    int idx1 = 0, idx2 = 0;
-    int stride1 = 1, stride2 = 1;
-    int linear_idx = i;
-
-    for (int j = max(ndim1, ndim2) - 1; j >= 0; j--) {
-        int dim1 = j < ndim1 ? shape1[ndim1 - 1 - j] : 1;
-        int dim2 = j < ndim2 ? shape2[ndim2 - 1 - j] : 1;
-        int broadcasted_dim = broadcasted_shape[j];
-
-        int pos = linear_idx % broadcasted_dim;
-        linear_idx /= broadcasted_dim;
-
-        if (dim1 > 1) {
-            idx1 += pos * stride1;
-        }
-        if (dim2 > 1) {
-            idx2 += pos * stride2;
-        }
-
-        stride1 *= dim1;
-        stride2 *= dim2;
+    
+    int index1 = 0, index2 = 0;
+    int linear_index = i;
+    for (int j = max_ndim - 1; j >= 0; j--) {
+        int pos = linear_index % broadcasted_shape[j];
+        linear_index /= broadcasted_shape[j];
+        if (strides1[j] != 0) index1 += pos * strides1[j];
+        if (strides2[j] != 0) index2 += pos * strides2[j];
     }
-
-    result_data[i] = data1[idx1] + data2[idx2];
+    result_data[i] = data1[index1] + data2[index2];
 }
 
 __host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape) {
-    int size = tensor1->size;
+    int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
 
-    // Copy the shapes to device memory
-    int* d_shape1;
-    int* d_shape2;
+    int* strides1 = (int*)malloc(max_ndim * sizeof(int));
+    int* strides2 = (int*)malloc(max_ndim * sizeof(int));
+    if (strides1 == NULL || strides2 == NULL) {
+        fprintf(stderr, "Memory allocation failed\n");
+        exit(1);
+    }
+
+    int stride1 = 1, stride2 = 1;
+    for (int i = max_ndim - 1; i >= 0; i--) {
+        int dim1 = i < tensor1->ndim ? tensor1->shape[tensor1->ndim - max_ndim + i] : 1;
+        int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
+        strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
+        strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
+        stride1 *= broadcasted_shape[i];
+        stride2 *= broadcasted_shape[i];
+    }
+    
     int* d_broadcasted_shape;
-    int ndim1 = tensor1->ndim;
-    int ndim2 = tensor2->ndim;
+    int* d_strides1;
+    int* d_strides2;
 
-    cudaMalloc((void**)&d_shape1, ndim1 * sizeof(int));
-    cudaMalloc((void**)&d_shape2, ndim2 * sizeof(int));
-    cudaMalloc((void**)&d_broadcasted_shape, max(ndim1, ndim2) * sizeof(int));
+    cudaMalloc((void**)&d_broadcasted_shape, max_ndim * sizeof(int));
+    cudaMemcpy(d_broadcasted_shape, broadcasted_shape, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
 
-    cudaMemcpy(d_shape1, tensor1->shape, ndim1 * sizeof(int), cudaMemcpyHostToDevice);
-    cudaMemcpy(d_shape2, tensor2->shape, ndim2 * sizeof(int), cudaMemcpyHostToDevice);
-    cudaMemcpy(d_broadcasted_shape, broadcasted_shape, max(ndim1, ndim2) * sizeof(int), cudaMemcpyHostToDevice);
+    cudaMalloc((void**)&d_strides1, max_ndim * sizeof(int));
+    cudaMemcpy(d_strides1, strides1, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
 
-    int number_of_blocks = (size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
-    add_broadcasted_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor1->data, tensor2->data, result_data, d_shape1, d_shape2, d_broadcasted_shape, ndim1, ndim2, size);
+    cudaMalloc((void**)&d_strides2, max_ndim * sizeof(int));
+    cudaMemcpy(d_strides2, strides2, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
+
+    int number_of_blocks = (tensor1->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    add_broadcasted_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor1->data, tensor2->data, result_data, d_broadcasted_shape, d_strides1, d_strides2, max_ndim, tensor1->size);
 
     cudaError_t error = cudaGetLastError();
     if (error != cudaSuccess) {
@@ -116,9 +114,6 @@ __host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, floa
     }
 
     cudaDeviceSynchronize();
-
-    cudaFree(d_shape1);
-    cudaFree(d_shape2);
     cudaFree(d_broadcasted_shape);
 }
 
diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h
index aeafc4e..e333ea7 100644
--- a/norch/csrc/cuda.h
+++ b/norch/csrc/cuda.h
@@ -7,7 +7,7 @@
     __global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
     __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);
 
-    __global__ void add_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* shape1, int* shape2, int* broadcasted_shape, int ndim1, int ndim2, int size);
+    __global__ void add_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size);
     __host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape);
 
     __global__ void sum_tensor_cuda_kernel(float* data, float* result_data);
diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp
index 4118e95..26b6c5a 100644
--- a/norch/csrc/tensor.cpp
+++ b/norch/csrc/tensor.cpp
@@ -150,11 +150,10 @@ extern "C" {
         }
 
         if (strcmp(tensor1->device, "cuda") == 0) {
-
             float* result_data;
             cudaMalloc((void **)&result_data, tensor1->size * sizeof(float));
-            add_broadcasted_tensor_cuda(tensor1, tensor2, result_data);
-            return create_tensor(result_data, shape, ndim, device);
+            add_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape);
+            return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
         } 
         else {
             float* result_data = (float*)malloc(tensor1->size * sizeof(float));
diff --git a/norch/libtensor.so b/norch/libtensor.so
index dd54141..c0f88c2 100755
Binary files a/norch/libtensor.so and b/norch/libtensor.so differ
diff --git a/tests/test_autograd.py b/tests/test_autograd.py
index ead6e1a..d83f5c4 100644
--- a/tests/test_autograd.py
+++ b/tests/test_autograd.py
@@ -10,6 +10,9 @@ class TestTensorAutograd(unittest.TestCase):
         if self.device is None or self.device != 'cuda':
             self.device = 'cpu'
         
+        print(f"Running tests on: {self.device}")
+
+        
     def test_addition(self):
         """
         Test autograd from addition two tensors: tensor1 + tensor2
diff --git a/tests/test_nn.py b/tests/test_nn.py
index 224c214..dca9c0e 100644
--- a/tests/test_nn.py
+++ b/tests/test_nn.py
@@ -11,6 +11,8 @@ class TestNNModuleLoss(unittest.TestCase):
         if self.device is None or self.device != 'cuda':
             self.device = 'cpu'
 
+        print(f"Running tests on: {self.device}")
+
     def test_mse_loss(self):
         """
         Test the MSELoss
diff --git a/tests/test_operations.py b/tests/test_operations.py
index a48f7c0..03f986c 100644
--- a/tests/test_operations.py
+++ b/tests/test_operations.py
@@ -12,6 +12,8 @@ class TestTensorOperations(unittest.TestCase):
         if self.device is None or self.device != 'cuda':
             self.device = 'cpu'
 
+        print(f"Running tests on: {self.device}")
+
     def test_creation_and_conversion(self):
         """
         Test creation and convertion of norch tensor to pytorch
@@ -50,6 +52,12 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
+        norch_result = norch_tensor2 + norch_tensor1
+        torch_expected = torch_tensor2 + torch_tensor1
+
+        self.assertTrue(utils.compare_torch(torch_result, torch_expected))
+
+
 
     def test_subtraction(self):
         """