Do not wrap array in a torch tensor if already a tensor

ami-iit · Jan 25, 2024 · 9095103 · 9095103
1 parent f308a52
commit 9095103
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/adam/pytorch/torch_like.py b/src/adam/pytorch/torch_like.py
@@ -59,6 +59,8 @@ def __matmul__(self, other: Union["TorchLike", ntp.ArrayLike]) -> "TorchLike":
 
         if type(self) is type(other):
             return TorchLike(self.array @ other.array)
+        if isinstance(other, torch.Tensor):
+            return TorchLike(self.array @ other)
         else:
             return TorchLike(self.array @ torch.tensor(other))