amend

pytorch · Oct 10, 2023 · 85f596d · 85f596d
1 parent 9ae41b8
commit 85f596d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tensordict/memmap.py b/tensordict/memmap.py
@@ -731,7 +731,8 @@ def to(
                 out.device = device
             return out
 
-        self.device = device
+        out = self.clone()
+        out.device = device
         return self
 
     def unbind(self, dim: int) -> tuple[torch.Tensor, ...]:
@@ -845,7 +846,7 @@ def _cat(
 def _where(condition, input, other):
     device = input.device
     if device != torch.device("cpu"):
-        input = input.to("cpu").as_tensor().to(device, non_blocking=True)
+        input = input.to("cpu").to(device, non_blocking=True)
     else:
         input = input.as_tensor()
     if condition.device != device or (isinstance(other, (MemmapTensor, torch.Tensor)) and other.device != device):