aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix send to device

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-23 09:46:35 -03:00
parent a072b57b50
commit 065e913b2d
3 changed files with 6 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
norch/nn/module.py
View file

@ -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()

2
norch/tensor.py
View file

@ -52,6 +52,8 @@ class Tensor:
self.ndim_ctype,
self.device_ctype
)
self.to(device)
else:
self.tensor = None,

7
tests/test_autograd.py
View file

@ -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()