diff --git a/tests/test_autograd.py b/tests/test_autograd.py index 9065471..4de2838 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -980,7 +980,7 @@ class TestTensorAutograd(unittest.TestCase): torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True, device=self.device) torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True, device=self.device) - + torch_result_transpose_matmul = (torch_tensor1.T @ torch_tensor2).sum() torch_result_transpose_matmul.backward() torch_tensor_grad_transpose_matmul1 = torch_tensor1.grad diff --git a/tests/test_nn.py b/tests/test_nn.py index 4896506..532f5b6 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -26,11 +26,8 @@ class TestNNModuleLoss(unittest.TestCase): loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) loss_torch_result = utils.to_torch(loss_norch).to(self.device) - predictions_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 4]]) - labels_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 3]]) - - predictions_torch.to(self.device) - labels_torch.to(self.device) + predictions_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 4]], device=self.device) + labels_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 3]], device=self.device) loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) @@ -42,11 +39,8 @@ class TestNNModuleLoss(unittest.TestCase): loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) loss_torch_result = utils.to_torch(loss_norch).to(self.device) - predictions_torch = torch.tensor([1.1, 2, 3, 4]) - labels_torch = torch.tensor([4, 3, 2.1, 1]) - - predictions_torch.to(self.device) - labels_torch.to(self.device) + predictions_torch = torch.tensor([1.1, 2, 3, 4], device=self.device) + labels_torch = torch.tensor([4, 3, 2.1, 1], device=self.device) loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) @@ -66,11 +60,8 @@ class TestNNModuleLoss(unittest.TestCase): loss_torch_result = utils.to_torch(loss_norch).to(self.device) - predictions_torch = torch.tensor([2.0, 1.0, 0.1]) - labels_torch = torch.tensor(0) - - predictions_torch.to(self.device) - labels_torch.to(self.device) + predictions_torch = torch.tensor([2.0, 1.0, 0.1], device=self.device) + labels_torch = torch.tensor(0, device=self.device) loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) @@ -83,11 +74,8 @@ class TestNNModuleLoss(unittest.TestCase): loss_torch_result = utils.to_torch(loss_norch).to(self.device) - predictions_torch = torch.tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]]) - labels_torch = torch.tensor([2, 1]) - - predictions_torch.to(self.device) - labels_torch.to(self.device) + predictions_torch = torch.tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]], device=self.device) + labels_torch = torch.tensor([2, 1], device=self.device) loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) @@ -99,11 +87,8 @@ class TestNNModuleLoss(unittest.TestCase): loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) loss_torch_result = utils.to_torch(loss_norch).to(self.device) - predictions_torch = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]) - labels_torch = torch.tensor([1, 2]) - - predictions_torch.to(self.device) - labels_torch.to(self.device) + predictions_torch = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], device=self.device) + labels_torch = torch.tensor([1, 2], device=self.device) loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) @@ -131,11 +116,8 @@ class TestNNModuleLoss(unittest.TestCase): loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) loss_torch_result = utils.to_torch(loss_norch).to(self.device) - predictions_torch = torch.tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]]) - labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]]) - - predictions_torch.to(self.device) - labels_torch.to(self.device) + predictions_torch = torch.tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]], device=self.device) + labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]], device=self.device) loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) @@ -161,9 +143,7 @@ class TestNNModuleActivationFn(unittest.TestCase): sigmoid_norch = sigmoid_fn_norch.forward(x) sigmoid_torch_result = utils.to_torch(sigmoid_norch).to(self.device) - x = torch.tensor([[1, 2, 3]]) - - x.to(self.device) + x = torch.tensor([[1, 2, 3]], device=self.device) sigmoid_torch_expected = sigmoid_fn_torch.forward(x) @@ -174,9 +154,7 @@ class TestNNModuleActivationFn(unittest.TestCase): sigmoid_norch = sigmoid_fn_norch.forward(x) sigmoid_torch_result = utils.to_torch(sigmoid_norch).to(self.device) - x = torch.tensor([-1, 2, -3]) - - x.to(self.device) + x = torch.tensor([-1, 2, -3], device=self.device) sigmoid_torch_expected = sigmoid_fn_torch.forward(x) @@ -187,9 +165,7 @@ class TestNNModuleActivationFn(unittest.TestCase): sigmoid_norch = sigmoid_fn_norch.forward(x) sigmoid_torch_result = utils.to_torch(sigmoid_norch).to(self.device) - x = torch.tensor([0, 0, 0]) - - x.to(self.device) + x = torch.tensor([0, 0, 0], device=self.device) sigmoid_torch_expected = sigmoid_fn_torch.forward(x) @@ -205,9 +181,9 @@ class TestNNModuleActivationFn(unittest.TestCase): # Define the input tensors for different test cases test_cases = [ - (norch.Tensor([[[1., 2, 3], [4, 5, 6]]]), torch.tensor([[[1., 2, 3], [4, 5, 6]]])), - (norch.Tensor([[[1., -1, 0], [2, -2, 0]]]), torch.tensor([[[1., -1, 0], [2, -2, 0]]])), - (norch.Tensor([[[0., 0, 0], [0, 0, 0]]]), torch.tensor([[[0., 0, 0], [0, 0, 0]]])) + (norch.Tensor([[[1., 2, 3], [4, 5, 6]]]).to(self.device), torch.tensor([[[1., 2, 3], [4, 5, 6]]], device=self.device)), + (norch.Tensor([[[1., -1, 0], [2, -2, 0]]]).to(self.device), torch.tensor([[[1., -1, 0], [2, -2, 0]]], device=self.device)), + (norch.Tensor([[[0., 0, 0], [0, 0, 0]]]).to(self.device), torch.tensor([[[0., 0, 0], [0, 0, 0]]], device=self.device)) ] for dim in axes: @@ -215,12 +191,7 @@ class TestNNModuleActivationFn(unittest.TestCase): softmax_fn_torch = torch.nn.Softmax(dim=dim) for norch_input, torch_input in test_cases: - # Move tensors to the correct device - norch_input = norch_input.to(self.device) - torch_input = torch_input - torch_input.to(self.device) - # Forward pass using norch softmax_norch = softmax_fn_norch.forward(norch_input) softmax_torch_result = utils.to_torch(softmax_norch).to(self.device)