fixed bug in copying to a new device in torch repr of mlnetwork

anty-filidor · Jul 5, 2024 · 128d426 · 128d426
1 parent 4ca0048
commit 128d426
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/network_diffusion/mln/mlnetwork_torch.py b/network_diffusion/mln/mlnetwork_torch.py
@@ -156,8 +156,8 @@ def device(self) -> str:
     @device.setter
     def device(self, new_device: str) -> None:
         """Copy tensor-members of the object into a `new_device`."""
-        self.adjacency_tensor.to(new_device)
-        self.nodes_mask.to(new_device)
+        self.adjacency_tensor = self.adjacency_tensor.to(new_device)
+        self.nodes_mask = self.nodes_mask.to(new_device)
 
     @classmethod
     def from_mln(