amend

pytorch · Oct 10, 2023 · b78eccf · b78eccf
1 parent 86ac8fa
commit b78eccf
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
@@ -1764,7 +1764,10 @@ def map(
     @cache  # noqa: B019
     def _add_batch_dim(self, *, in_dim, vmap_level):
         if self.is_memmap():
-            td = self.cpu()
+            if self.device.type != "cpu":
+                td = self.cpu()
+            else:
+                td = self.as_tensor()
         else:
             td = self
         out = TensorDict(