amend

pytorch · Oct 10, 2023 · 7665531 · 7665531
1 parent 2fb64be
commit 7665531
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tensordict/memmap.py b/tensordict/memmap.py
@@ -843,7 +843,12 @@ def _cat(
 
 
 def _where(condition, input, other):
-    return torch.where(condition=condition, input=input.as_tensor(), other=other)
+    device = input.device
+    if device != torch.device("cpu"):
+        input = input.to("cpu").as_tensor().to(device)
+    else:
+        input = input.as_tensor()
+    return torch.where(condition=condition, input=input, other=other)
 
 
 implements_for_memmap(torch.where)(_where)