Merge branch 'sampler-doc' of github.com:pytorch/rl into sampler-doc

pytorch · Feb 7, 2024 · c203b12 · c203b12
2 parents a8a6537 + be46f16
commit c203b12
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/torchrl/data/replay_buffers/samplers.py b/torchrl/data/replay_buffers/samplers.py
@@ -435,8 +435,10 @@ def update_priority(
                 indexed elements.
 
         """
-        priority = torch.as_tensor(priority, dtype=torch.long)
-        index = torch.as_tensor(index, dtype=torch.long)
+        priority = torch.as_tensor(
+            priority, dtype=torch.long, device=torch.device("cpu")
+        )
+        index = torch.as_tensor(index, dtype=torch.long, device=torch.device("cpu"))
         # we need to reshape priority if it has more than one elements or if it has
         # a different shape than index
         if priority.numel() > 1 and priority.shape != index.shape: