From 065e913b2df047349a6c170920e3f0051b07ee5b Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Thu, 23 May 2024 09:46:35 -0300 Subject: [PATCH] fix send to device --- norch/nn/module.py | 2 ++ norch/tensor.py | 2 ++ tests/test_autograd.py | 7 ++----- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/norch/nn/module.py b/norch/nn/module.py index cebd86f..eaabd9a 100644 --- a/norch/nn/module.py +++ b/norch/nn/module.py @@ -53,6 +53,8 @@ class Module(ABC): def to(self, device): for _, _, parameter in self.parameters(): parameter.to(device) + + return self def state_dict(self): state = OrderedDict() diff --git a/norch/tensor.py b/norch/tensor.py index 8063cba..e660df9 100644 --- a/norch/tensor.py +++ b/norch/tensor.py @@ -52,6 +52,8 @@ class Tensor: self.ndim_ctype, self.device_ctype ) + + self.to(device) else: self.tensor = None, diff --git a/tests/test_autograd.py b/tests/test_autograd.py index 6793816..823ba1b 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -24,11 +24,8 @@ class TestTensorAutograd(unittest.TestCase): norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device) norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device) - 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.to(self.device) - torch_tensor2.to(self.device) + torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True, device=self.device) + torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True, device=self.device) torch_result = (torch_tensor1 + torch_tensor2).sum() torch_result.backward()