[BugFix] Remove split_trajectories's reference to `("next", "done")…

…`. (#2094)
pytorch · Apr 22, 2024 · 32533a5 · 32533a5
1 parent 68ef60b
commit 32533a5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/collectors/utils.py b/torchrl/collectors/utils.py
@@ -179,7 +179,7 @@ def split_trajectories(
             torch.ones(
                 out_split.shape,
                 dtype=torch.bool,
-                device=out_split.get(("next", "done")).device,
+                device=out_split.device,
             ),
         )
     if len(out_splits) > 1: