amend

pytorch · Oct 10, 2023 · 9ae41b8 · 9ae41b8
1 parent c47ed6e
commit 9ae41b8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensordict/memmap.py b/tensordict/memmap.py
@@ -848,7 +848,7 @@ def _where(condition, input, other):
         input = input.to("cpu").as_tensor().to(device, non_blocking=True)
     else:
         input = input.as_tensor()
-    if condition.device != device or other.device != device:
+    if condition.device != device or (isinstance(other, (MemmapTensor, torch.Tensor)) and other.device != device):
         raise ValueError(f"{condition.device}, {device}, {other.device}")
     return torch.where(condition=condition, input=input, other=other)