fix: TensorArray .to() moves tensor now to device

MIC-DKFZ · Jul 2, 2024 · ede69e2 · ede69e2
1 parent 19abf27
commit ede69e2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/patchly/array_like.py b/patchly/array_like.py
@@ -227,4 +227,4 @@ def argmax(self, axis=None):
         return TensorArray(torch.argmax(self.data, axis), self.device)
 
     def to(self, device):
-        return TensorArray(self.data, self.device)
+        return TensorArray(self.data, device)