[BugFix] Fix storage device (#1650)

pytorch · Oct 24, 2023 · 105e861 · 105e861
1 parent b7d148b
commit 105e861
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/torchrl/data/replay_buffers/storages.py b/torchrl/data/replay_buffers/storages.py
@@ -372,7 +372,9 @@ def set(  # noqa: F811
                 self._init(data)
         if not isinstance(cursor, (*INT_CLASSES, slice)):
             if not isinstance(cursor, torch.Tensor):
-                cursor = torch.tensor(cursor)
+                cursor = torch.tensor(cursor, dtype=torch.long, device=self.device)
+            elif cursor.dtype != torch.long:
+                cursor = cursor.to(dtype=torch.long, device=self.device)
             if len(cursor) > len(self._storage):
                 warnings.warn(
                     "A cursor of length superior to the storage capacity was provided. "