diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc
index af582eb..1471e6d 100644
Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ
diff --git a/test.py b/test.py
index 7c26de2..dcbda9f 100644
--- a/test.py
+++ b/test.py
@@ -22,7 +22,7 @@ if __name__ == "__main__":
     import psutil
     from norch.utils import utils
 
-
+    """
 
     a = norch.Tensor([
         [[1.234, 2.123], [3.635, 4.456], [5.678, 6.789]],
@@ -49,7 +49,7 @@ if __name__ == "__main__":
 
     print(utils.torch_compare(a, b))
 
-    exit()
+    exit()"""
 
     """
 
@@ -92,7 +92,7 @@ if __name__ == "__main__":
 
     print(a.grad)"""
 
-    import norch.nn as nn
+    """import norch.nn as nn
 
     cpu_percent = psutil.cpu_percent(interval=1)
     print(f"CPU Usage: {cpu_percent}%")
@@ -139,7 +139,7 @@ if __name__ == "__main__":
         #print(loss)
 
     fim = time.time()
-    print(fim - ini)
+    print(fim - ini)"""
 
 
     #### testar transpose axes!!!! make it contiguous
@@ -192,15 +192,14 @@ if __name__ == "__main__":
 
     #d = b-c
 
-    """a = norch.Tensor([[1, 2], [1, 2], [1, 2]], requires_grad=True)#.to("cuda")
-    b = norch.Tensor([[1, 400, 3], [1, 2, 3]], requires_grad=True)
-
-    c = (a@b).reshape([9])
+    a = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True)#.to("cuda")
+    b = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True)
+    t = b.reshape([2,4])
+    c = (t @ a)
     d = c.sum()
     d.backward()
 
-    print(a.grad)"""
-    
+    print(a.grad)
     """#print(a)
     N = 10
     a = norch.Tensor([[1 for _ in range(N)] for _ in range(N)])
diff --git a/tests/test_autograd.py b/tests/test_autograd.py
index 0615b4f..2ac0d92 100644
--- a/tests/test_autograd.py
+++ b/tests/test_autograd.py
@@ -6,21 +6,23 @@ import os
 
 class TestTensorAutograd(unittest.TestCase):
     def setUp(self):
-        self.device = os.environ.get('device', 'cpu')
+        self.device = os.environ.get('device')
+        if self.device is None or self.device != 'cuda':
+            self.device = 'cpu'
         
     def test_addition(self):
         """
         Test autograd from addition two tensors: tensor1 + tensor2
         """
-        norch_tensor1 = norch.Tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
-        norch_tensor2 = norch.Tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True)
+        norch_tensor1 = norch.Tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        norch_tensor2 = norch.Tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
         norch_result = (norch_tensor1 + norch_tensor2).sum()
         norch_result.backward()
         norch_tensor1_grad = utils.to_torch(norch_tensor1.grad)
         norch_tensor2_grad = utils.to_torch(norch_tensor2.grad)
 
-        torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
-        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True)
+        torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
         torch_result = (torch_tensor1 + torch_tensor2).sum()
         torch_result.backward()
         torch_tensor1_grad = torch_tensor1.grad
@@ -34,15 +36,15 @@ class TestTensorAutograd(unittest.TestCase):
         """
         Test autograd from subtraction two tensors: tensor1 - tensor2
         """
-        norch_tensor1_sub = norch.Tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
-        norch_tensor2_sub = norch.Tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True)
+        norch_tensor1_sub = norch.Tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        norch_tensor2_sub = norch.Tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
         norch_result_sub = (norch_tensor1_sub - norch_tensor2_sub).sum()
         norch_result_sub.backward()
         norch_tensor1_grad_sub = utils.to_torch(norch_tensor1_sub.grad)
         norch_tensor2_grad_sub = utils.to_torch(norch_tensor2_sub.grad)
 
-        torch_tensor1_sub = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
-        torch_tensor2_sub = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True)
+        torch_tensor1_sub = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor2_sub = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
         torch_result_sub = (torch_tensor1_sub - torch_tensor2_sub).sum()
         torch_result_sub.backward()
         torch_tensor1_grad_sub = torch_tensor1_sub.grad
@@ -55,15 +57,15 @@ class TestTensorAutograd(unittest.TestCase):
         """
         Test autograd from dividing two tensors: tensor1 / tensor2
         """
-        norch_tensor1_div = norch.Tensor([[[2, 5.1], [6, -8]], [[10, 12], [14, 16]]], requires_grad=True)
-        norch_tensor2_div = norch.Tensor([[[1, 1], [2, 2.2]], [[3, 3], [4, 4]]], requires_grad=True)
+        norch_tensor1_div = norch.Tensor([[[2, 5.1], [6, -8]], [[10, 12], [14, 16]]], requires_grad=True).to(self.device)
+        norch_tensor2_div = norch.Tensor([[[1, 1], [2, 2.2]], [[3, 3], [4, 4]]], requires_grad=True).to(self.device)
         norch_result_div = (norch_tensor1_div / norch_tensor2_div).sum()
         norch_result_div.backward()
         norch_tensor1_grad_div = utils.to_torch(norch_tensor1_div.grad)
         norch_tensor2_grad_div = utils.to_torch(norch_tensor2_div.grad)
 
-        torch_tensor1_div = torch.tensor([[[2, 5.1], [6, -8]], [[10, 12], [14, 16]]], requires_grad=True)
-        torch_tensor2_div = torch.tensor([[[1, 1], [2, 2.2]], [[3, 3], [4, 4]]], requires_grad=True)
+        torch_tensor1_div = torch.tensor([[[2, 5.1], [6, -8]], [[10, 12], [14, 16]]], requires_grad=True).to(self.device)
+        torch_tensor2_div = torch.tensor([[[1, 1], [2, 2.2]], [[3, 3], [4, 4]]], requires_grad=True).to(self.device)
         torch_result_div = (torch_tensor1_div / torch_tensor2_div).sum()
         torch_result_div.backward()
         torch_tensor1_grad_div = torch_tensor1_div.grad
@@ -77,13 +79,13 @@ class TestTensorAutograd(unittest.TestCase):
         """
         Test autograd from dividing tensor by scalar: tensor / scalar
         """
-        norch_tensor_div_scalar = norch.Tensor([[[2, 4.7], [6, 8]], [[10, 12], [14, 16]]], requires_grad=True)
+        norch_tensor_div_scalar = norch.Tensor([[[2, 4.7], [6, 8]], [[10, 12], [14, 16]]], requires_grad=True).to(self.device)
         scalar = 2
         norch_result_div_scalar = (norch_tensor_div_scalar / scalar).sum()
         norch_result_div_scalar.backward()
         norch_tensor_grad_div_scalar = utils.to_torch(norch_tensor_div_scalar.grad)
 
-        torch_tensor_div_scalar = torch.tensor([[[2, 4.7], [6, 8]], [[10, 12], [14, 16]]], requires_grad=True)
+        torch_tensor_div_scalar = torch.tensor([[[2, 4.7], [6, 8]], [[10, 12], [14, 16]]], requires_grad=True).to(self.device)
         torch_result_div_scalar = (torch_tensor_div_scalar / scalar).sum()
         torch_result_div_scalar.backward()
         torch_tensor_grad_div_scalar = torch_tensor_div_scalar.grad
@@ -96,12 +98,12 @@ class TestTensorAutograd(unittest.TestCase):
         Test autograd from dividing scalar by tensor: scalar / tensor
         """
         scalar = 2
-        norch_tensor_scalar_div = norch.Tensor([[[1, 2.23], [3, 4]], [[5, 6], [7, 8]]], requires_grad=True)
+        norch_tensor_scalar_div = norch.Tensor([[[1, 2.23], [3, 4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         norch_result_scalar_div = (scalar / norch_tensor_scalar_div).sum()
         norch_result_scalar_div.backward()
         norch_tensor_grad_scalar_div = utils.to_torch(norch_tensor_scalar_div.grad)
 
-        torch_tensor_scalar_div = torch.tensor([[[1, 2.23], [3, 4]], [[5, 6], [7, 8]]], requires_grad=True)
+        torch_tensor_scalar_div = torch.tensor([[[1, 2.23], [3, 4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         torch_result_scalar_div = (scalar / torch_tensor_scalar_div).sum()
         torch_result_scalar_div.backward()
         torch_tensor_grad_scalar_div = torch_tensor_scalar_div.grad
@@ -114,12 +116,12 @@ class TestTensorAutograd(unittest.TestCase):
         Test autograd from scalar raised to tensor: scalar ** tensor
         """
         scalar = 2
-        norch_tensor_power_st = norch.Tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True)
+        norch_tensor_power_st = norch.Tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
         norch_result_power_st = (scalar ** norch_tensor_power_st).sum()
         norch_result_power_st.backward()
         norch_tensor_grad_power_st = utils.to_torch(norch_tensor_power_st.grad)
 
-        torch_tensor_power_st = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True)
+        torch_tensor_power_st = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
         torch_result_power_st = (scalar ** torch_tensor_power_st).sum()
         torch_result_power_st.backward()
         torch_tensor_grad_power_st = torch_tensor_power_st.grad
@@ -132,12 +134,12 @@ class TestTensorAutograd(unittest.TestCase):
         Test autograd from tensor raised to scalar: tensor ** scalar
         """
         scalar = 2
-        norch_tensor_power_ts = norch.Tensor([[[2, 3], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True)
+        norch_tensor_power_ts = norch.Tensor([[[2, 3], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
         norch_result_power_ts = (norch_tensor_power_ts ** scalar).sum()
         norch_result_power_ts.backward()
         norch_tensor_grad_power_ts = utils.to_torch(norch_tensor_power_ts.grad)
 
-        torch_tensor_power_ts = torch.tensor([[[2, 3], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True)
+        torch_tensor_power_ts = torch.tensor([[[2, 3], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
         torch_result_power_ts = (torch_tensor_power_ts ** scalar).sum()
         torch_result_power_ts.backward()
         torch_tensor_grad_power_ts = torch_tensor_power_ts.grad
@@ -148,15 +150,15 @@ class TestTensorAutograd(unittest.TestCase):
         """
         Test autograd from matrix multiplication: matmul(tensor1, tensor2)
         """
-        norch_tensor1_matmul = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
-        norch_tensor2_matmul = norch.Tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True)
+        norch_tensor1_matmul = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        norch_tensor2_matmul = norch.Tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
         norch_result_matmul = (norch_tensor1_matmul @ norch_tensor2_matmul).sum()
         norch_result_matmul.backward()
         norch_tensor1_grad_matmul = utils.to_torch(norch_tensor1_matmul.grad)
         norch_tensor2_grad_matmul = utils.to_torch(norch_tensor2_matmul.grad)
 
-        torch_tensor1_matmul = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
-        torch_tensor2_matmul = torch.tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True)
+        torch_tensor1_matmul = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor2_matmul = torch.tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
         torch_result_matmul = (torch_tensor1_matmul @ torch_tensor2_matmul).sum()
         torch_result_matmul.backward()
         torch_tensor1_grad_matmul = torch_tensor1_matmul.grad
@@ -171,12 +173,12 @@ class TestTensorAutograd(unittest.TestCase):
         Test autograd from elementwise multiplication with scalar: scalar * tensor
         """
         scalar = 2
-        norch_tensor_elemwise_mul_scalar = norch.Tensor([[[1.1, 2], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        norch_tensor_elemwise_mul_scalar = norch.Tensor([[[1.1, 2], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         norch_result_elemwise_mul_scalar = (scalar * norch_tensor_elemwise_mul_scalar).sum()
         norch_result_elemwise_mul_scalar.backward()
         norch_tensor_grad_elemwise_mul_scalar = utils.to_torch(norch_tensor_elemwise_mul_scalar.grad)
 
-        torch_tensor_elemwise_mul_scalar = torch.tensor([[[1.1, 2], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        torch_tensor_elemwise_mul_scalar = torch.tensor([[[1.1, 2], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         torch_result_elemwise_mul_scalar = (scalar * torch_tensor_elemwise_mul_scalar).sum()
         torch_result_elemwise_mul_scalar.backward()
         torch_tensor_grad_elemwise_mul_scalar = torch_tensor_elemwise_mul_scalar.grad
@@ -188,15 +190,15 @@ class TestTensorAutograd(unittest.TestCase):
         """
         Test autograd from elementwise multiplication between two tensors: tensor1 * tensor2
         """
-        norch_tensor1_elemwise_mul = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
-        norch_tensor2_elemwise_mul = norch.Tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True)
+        norch_tensor1_elemwise_mul = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        norch_tensor2_elemwise_mul = norch.Tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
         norch_result_elemwise_mul = (norch_tensor1_elemwise_mul * norch_tensor2_elemwise_mul).sum()
         norch_result_elemwise_mul.backward()
         norch_tensor1_grad_elemwise_mul = utils.to_torch(norch_tensor1_elemwise_mul.grad)
         norch_tensor2_grad_elemwise_mul = utils.to_torch(norch_tensor2_elemwise_mul.grad)
 
-        torch_tensor1_elemwise_mul = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
-        torch_tensor2_elemwise_mul = torch.tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True)
+        torch_tensor1_elemwise_mul = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor2_elemwise_mul = torch.tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
         torch_result_elemwise_mul = (torch_tensor1_elemwise_mul * torch_tensor2_elemwise_mul).sum()
         torch_result_elemwise_mul.backward()
         torch_tensor1_grad_elemwise_mul = torch_tensor1_elemwise_mul.grad
@@ -209,13 +211,13 @@ class TestTensorAutograd(unittest.TestCase):
         """
         Test autograd from reshaping a tensor: tensor.reshape(shape)
         """
-        norch_tensor_reshape = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        norch_tensor_reshape = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         new_shape = [2, 4]
         norch_result_reshape = norch_tensor_reshape.reshape(new_shape).sum()
         norch_result_reshape.backward()
         norch_tensor_grad_reshape = utils.to_torch(norch_tensor_reshape.grad)
 
-        torch_tensor_reshape = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        torch_tensor_reshape = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         torch_result_reshape = torch_tensor_reshape.reshape(new_shape).sum()
         torch_result_reshape.backward()
         torch_tensor_grad_reshape = torch_tensor_reshape.grad
@@ -227,13 +229,13 @@ class TestTensorAutograd(unittest.TestCase):
         """
         Test autograd from transposing a tensor with specific axes: tensor.transpose(axis1, axis2)
         """
-        norch_tensor_transpose = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        norch_tensor_transpose = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         axis1, axis2 = 0, 2
         norch_result_transpose = norch_tensor_transpose.transpose(axis1, axis2).sum()
         norch_result_transpose.backward()
         norch_tensor_grad_transpose = utils.to_torch(norch_tensor_transpose.grad)
 
-        torch_tensor_transpose = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        torch_tensor_transpose = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         torch_result_transpose = torch_tensor_transpose.transpose(axis1, axis2).sum()
         torch_result_transpose.backward()
         torch_tensor_grad_transpose = torch_tensor_transpose.grad
@@ -245,12 +247,12 @@ class TestTensorAutograd(unittest.TestCase):
         """
         Test autograd from transposing a tensor using .T attribute
         """
-        norch_tensor_T = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        norch_tensor_T = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         norch_result_T = norch_tensor_T.T.sum()
         norch_result_T.backward()
         norch_tensor_grad_T = utils.to_torch(norch_tensor_T.grad)
 
-        torch_tensor_T = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        torch_tensor_T = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         torch_result_T = torch_tensor_T.mT.sum()
         torch_result_T.backward()
         torch_tensor_grad_T = torch_tensor_T.grad
@@ -261,58 +263,82 @@ class TestTensorAutograd(unittest.TestCase):
         """
         Test autograd from reshaping a tensor then performing matrix multiplication: matmul(tensor1.reshape(shape), tensor2)
         """
-        norch_tensor_reshape_matmul = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True)
-        new_shape = [2, 4]
-        norch_result_reshape_matmul = (norch_tensor_reshape_matmul.reshape(new_shape) @ norch_tensor_reshape_matmul).sum()
-        norch_result_reshape_matmul.backward()
-        norch_tensor_grad_reshape_matmul = utils.to_torch(norch_tensor_reshape_matmul.grad)
+        norch_tensor1 = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True).to(self.device)
+        norch_tensor2 = norch.Tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], requires_grad=True).to(self.device)
 
-        torch_tensor_reshape_matmul = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True)
-        torch_result_reshape_matmul = torch.matmul(torch_tensor_reshape_matmul.reshape(new_shape), torch_tensor_reshape_matmul).sum()
+        new_shape = [2, 4]
+
+        norch_result_reshape_matmul = (norch_tensor1.reshape(new_shape) @ norch_tensor2).sum()
+        norch_result_reshape_matmul.backward()
+        norch_tensor_grad_reshape_matmul1 = utils.to_torch(norch_tensor1.grad)
+        norch_tensor_grad_reshape_matmul2 = utils.to_torch(norch_tensor2.grad)
+
+        torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)
+        torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)     
+        
+        torch_result_reshape_matmul = (torch_tensor1.reshape(new_shape) @ torch_tensor2).sum()
         torch_result_reshape_matmul.backward()
-        torch_tensor_grad_reshape_matmul = torch_tensor_reshape_matmul.grad
+        torch_tensor_grad_reshape_matmul1 = torch_tensor1.grad
+        torch_tensor_grad_reshape_matmul2 = torch_tensor2.grad
         
-        
-        print(norch_tensor_grad_reshape_matmul)
-        print(torch_tensor_grad_reshape_matmul)
-        print("\n\n\n\n@@")
-        
-        self.assertTrue(utils.compare_torch(norch_tensor_grad_reshape_matmul, torch_tensor_grad_reshape_matmul))
+        self.assertTrue(utils.compare_torch(norch_tensor_grad_reshape_matmul1, torch_tensor_grad_reshape_matmul1))
+        self.assertTrue(utils.compare_torch(norch_tensor_grad_reshape_matmul2, torch_tensor_grad_reshape_matmul2))
 
 
     def test_T_then_matmul(self):
         """
         Test autograd from transposing a tensor then performing matrix multiplication: matmul(tensor.T, tensor)
         """
-        norch_tensor_T_matmul = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True)
-        norch_result_T_matmul = (norch_tensor_T_matmul.T @ norch_tensor_T_matmul).sum()
+        norch_tensor1 = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True)
+        norch_tensor2 = norch.Tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], requires_grad=True)
+
+        norch_result_T_matmul = (norch_tensor1.T @ norch_tensor2).sum()
         norch_result_T_matmul.backward()
-        norch_tensor_grad_T_matmul = utils.to_torch(norch_tensor_T_matmul.grad)
+        norch_tensor_grad_T_matmul1 = utils.to_torch(norch_tensor1.grad)
+        norch_tensor_grad_T_matmult2 = utils.to_torch(norch_tensor2.grad)
 
-        torch_tensor_T_matmul = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True)
-        torch_result_T_matmul = torch.matmul(torch_tensor_T_matmul.T, torch_tensor_T_matmul).sum()
+        torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)
+        torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)     
+        
+        torch_result_T_matmul = (torch_tensor1.T @ torch_tensor2).sum()
         torch_result_T_matmul.backward()
-        torch_tensor_grad_T_matmul = torch_tensor_T_matmul.grad
-
-        self.assertTrue(utils.compare_torch(norch_tensor_grad_T_matmul, torch_tensor_grad_T_matmul))
+        torch_tensor_grad_T_matmul1 = torch_tensor1.grad
+        torch_tensor_grad_T_matmul2 = torch_tensor2.grad
+        
+        self.assertTrue(utils.compare_torch(norch_tensor_grad_T_matmul1, torch_tensor_grad_T_matmul1))
+        self.assertTrue(utils.compare_torch(norch_tensor_grad_T_matmult2, torch_tensor_grad_T_matmul2))
 
+    def todo(self):
+        """
+        The code has a problem on the following operation
+        tensor1.reshape(..) @ tensor1
+        print(tensor1.grad)
+        (also transpsoe and .T)
+        """
+        pass
 
     def test_transpose_axes_then_matmul(self):
         """
         Test autograd from transposing a tensor with specific axes then performing matrix multiplication: matmul(tensor.transpose(axis1, axis2), tensor)
         """
-        norch_tensor_transpose_matmul = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True)
-        axis1, axis2 = 0, 1
-        norch_result_transpose_matmul = (norch_tensor_transpose_matmul.transpose(axis1, axis2) @ norch_tensor_transpose_matmul).sum()
+        norch_tensor1 = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True).to(self.device)
+        norch_tensor2 = norch.Tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], requires_grad=True).to(self.device)
+
+        norch_result_transpose_matmul = (norch_tensor1.transpose(0, 1) @ norch_tensor2).sum()
         norch_result_transpose_matmul.backward()
-        norch_tensor_grad_transpose_matmul = utils.to_torch(norch_tensor_transpose_matmul.grad)
+        norch_tensor_grad_transpose_matmul1 = utils.to_torch(norch_tensor1.grad)
+        norch_tensor_grad_transpose_matmult2 = utils.to_torch(norch_tensor2.grad)
 
-        torch_tensor_transpose_matmul = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True)
-        torch_result_transpose_matmul = torch.matmul(torch_tensor_transpose_matmul.transpose(axis1, axis2), torch_tensor_transpose_matmul).sum()
+        torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)
+        torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)  
+        
+        torch_result_transpose_matmul = (torch_tensor1.T @ torch_tensor2).sum()
         torch_result_transpose_matmul.backward()
-        torch_tensor_grad_transpose_matmul = torch_tensor_transpose_matmul.grad
-
-        self.assertTrue(utils.compare_torch(norch_tensor_grad_transpose_matmul, torch_tensor_grad_transpose_matmul))
+        torch_tensor_grad_transpose_matmul1 = torch_tensor1.grad
+        torch_tensor_grad_transpose_matmul2 = torch_tensor2.grad
+        
+        self.assertTrue(utils.compare_torch(norch_tensor_grad_transpose_matmul1, torch_tensor_grad_transpose_matmul1))
+        self.assertTrue(utils.compare_torch(norch_tensor_grad_transpose_matmult2, torch_tensor_grad_transpose_matmul2))
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file
diff --git a/tests/test_operations.py b/tests/test_operations.py
index b01c231..6ded830 100644
--- a/tests/test_operations.py
+++ b/tests/test_operations.py
@@ -3,13 +3,20 @@ import norch
 from norch import utils
 import torch
 import sys
+import os
 
 class TestTensorOperations(unittest.TestCase):
+
+    def setUp(self):
+        self.device = os.environ.get('device')
+        if self.device is None or self.device != 'cuda':
+            self.device = 'cpu'
+
     def test_creation_and_conversion(self):
         """
         Test creation and convertion of norch tensor to pytorch
         """
-        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_tensor = utils.to_torch(norch_tensor)
         self.assertTrue(torch.is_tensor(torch_tensor))
 
@@ -17,13 +24,13 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test addition two tensors: tensor1 + tensor2
         """
-        norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
-        norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]])
+        norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
+        norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device)
         norch_result = norch_tensor1 + norch_tensor2
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
-        torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]])
+        torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
+        torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device)
         torch_expected = torch_tensor1 + torch_tensor2
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -32,13 +39,13 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test subtraction of two tensors: tensor1 - tensor2
         """
-        norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
-        norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]])
+        norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
+        norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device)
         norch_result = norch_tensor1 - norch_tensor2
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
-        torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]])
+        torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
+        torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device)
         torch_expected = torch_tensor1 - torch_tensor2
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -47,12 +54,12 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test division of a tensor by a scalar: tensor / scalar
         """
-        norch_tensor = norch.Tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]])
+        norch_tensor = norch.Tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device)
         scalar = 2
         norch_result = norch_tensor / scalar
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]])
+        torch_tensor = torch.tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device)
         torch_expected = torch_tensor / scalar
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -62,11 +69,11 @@ class TestTensorOperations(unittest.TestCase):
         Test scalar division by a tensor: scalar / tensor
         """
         scalar = 10
-        norch_tensor = norch.Tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]])
+        norch_tensor = norch.Tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device)
         norch_result = scalar / norch_tensor
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]])
+        torch_tensor = torch.tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device)
         torch_expected = scalar / torch_tensor
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -75,13 +82,13 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test matrix multiplication: tensor1 @ tensor2
         """
-        norch_tensor1 = norch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
-        norch_tensor2 = norch.Tensor([[[1, 0], [0, 1]], [[-1, 0], [0, -1]]])
+        norch_tensor1 = norch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device)
+        norch_tensor2 = norch.Tensor([[[1, 0], [0, 1]], [[-1, 0], [0, -1]]]).to(self.device)
         norch_result = norch_tensor1 @ norch_tensor2
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor1 = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
-        torch_tensor2 = torch.tensor([[[1, 0], [0, 1]], [[-1, 0], [0, -1]]])
+        torch_tensor1 = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device)
+        torch_tensor2 = torch.tensor([[[1, 0], [0, 1]], [[-1, 0], [0, -1]]]).to(self.device)
         torch_expected = torch_tensor1 @ torch_tensor2
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -90,12 +97,12 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test elementwise multiplication of a tensor by a scalar: tensor * scalar
         """
-        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         scalar = 2
         norch_result = norch_tensor * scalar
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch_tensor * scalar
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -104,13 +111,13 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test elementwise multiplication of two tensors: tensor1 * tensor2
         """
-        norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
-        norch_tensor2 = norch.Tensor([[[2, 2], [2, 2]], [[2, 2], [2, 2]]])
+        norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
+        norch_tensor2 = norch.Tensor([[[2, 2], [2, 2]], [[2, 2], [2, 2]]]).to(self.device)
         norch_result = norch_tensor1 * norch_tensor2
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
-        torch_tensor2 = torch.tensor([[[2, 2], [2, 2]], [[2, 2], [2, 2]]])
+        torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
+        torch_tensor2 = torch.tensor([[[2, 2], [2, 2]], [[2, 2], [2, 2]]]).to(self.device)
         torch_expected = torch_tensor1 * torch_tensor2
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -119,12 +126,12 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test reshaping of a tensor: tensor.reshape(shape)
         """
-        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         new_shape = [2, 4]
         norch_result = norch_tensor.reshape(new_shape)
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch_tensor.reshape(new_shape)
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -133,12 +140,12 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test transposition of a tensor: tensor.transpose(dim1, dim2)
         """
-        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         dim1, dim2 = 0, 2
         norch_result = norch_tensor.transpose(dim1, dim2)
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch_tensor.transpose(dim1, dim2)
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -147,11 +154,11 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test elementwise logarithm of a tensor: tensor.log()
         """
-        norch_tensor = norch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+        norch_tensor = norch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device)
         norch_result = norch_tensor.log()
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+        torch_tensor = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch.log(torch_tensor)
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -160,11 +167,11 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test summation of a tensor: tensor.sum()
         """
-        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         norch_result = norch_tensor.sum()
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch.sum(torch_tensor)
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -173,11 +180,11 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test transposition of a tensor: tensor.T
         """
-        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         norch_result = norch_tensor.T
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]])
+        torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch.transpose(torch_tensor, 0, 2)
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -186,14 +193,14 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test reshaping a tensor followed by matrix multiplication: (tensor.reshape(shape) @ other_tensor)
         """
-        norch_tensor = norch.Tensor([[1, 2], [3, -4], [5, 6], [7, 8]])
+        norch_tensor = norch.Tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device)
         new_shape = [2, 4]
         norch_reshaped = norch_tensor.reshape(new_shape)
         
         norch_result = norch_reshaped @ norch_tensor
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[1, 2], [3, -4], [5, 6], [7, 8]])
+        torch_tensor = torch.tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device)
         torch_expected = torch_tensor.reshape(new_shape) @ torch_tensor
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -203,12 +210,12 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test transposing a tensor followed by matrix multiplication: (tensor.transpose(dim1, dim2) @ other_tensor)
         """
-        norch_tensor = norch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+        norch_tensor = norch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device)
         dim1, dim2 = 0, 2
         norch_result = norch_tensor.transpose(dim1, dim2) @ norch_tensor
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+        torch_tensor = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch_tensor.transpose(dim1, dim2) @ torch_tensor
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -217,16 +224,16 @@ class TestTensorOperations(unittest.TestCase):
         """
         Test a combination of operations: (tensor.sum() + other_tensor) / scalar @ another_tensor followed by reshape
         """
-        norch_tensor1 = norch.Tensor([[[1., 2], [3, -4]], [[5, 6], [7, 8]]])
-        norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]])
+        norch_tensor1 = norch.Tensor([[[1., 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
+        norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device)
         scalar = 2
         new_shape = [2, 4]
         norch_result = ((norch_tensor1 + norch_tensor2) / scalar) @ norch_tensor1
         norch_result = norch_result.reshape(new_shape)
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor1 = torch.tensor([[[1., 2], [3, -4]], [[5, 6], [7, 8]]])
-        torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]])
+        torch_tensor1 = torch.tensor([[[1., 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
+        torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device)
         torch_expected = ((torch_tensor1 + torch_tensor2) / scalar) @ torch_tensor1
         torch_expected = torch_expected.reshape(new_shape)
 
@@ -237,11 +244,11 @@ class TestTensorOperations(unittest.TestCase):
         Test scalar power of a tensor: scalar ** tensor
         """
         scalar = 3
-        norch_tensor = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]])
+        norch_tensor = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         norch_result = scalar ** norch_tensor
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]])
+        torch_tensor = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = scalar ** torch_tensor
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
@@ -251,11 +258,11 @@ class TestTensorOperations(unittest.TestCase):
         Test tensor power of a scalar: tensor ** scalar
         """
         scalar = 3
-        norch_tensor = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]])
+        norch_tensor = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         norch_result = norch_tensor ** scalar
         torch_result = utils.to_torch(norch_result)
 
-        torch_tensor = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]])
+        torch_tensor = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch_tensor ** scalar
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))