[BugFix] Use traj_terminated for proper population of terminated key …

…in SliceSampler
pytorch · Feb 7, 2024 · 2d92eab · 2d92eab
1 parent 144f547
commit 2d92eab
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/data/replay_buffers/samplers.py b/torchrl/data/replay_buffers/samplers.py
@@ -917,7 +917,7 @@ def _sample_slices(
                 truncated[seq_length.cumsum(0) - 1] = 1
             traj_terminated = stop_idx[traj_idx] == start_idx[traj_idx] + seq_length - 1
             terminated = torch.zeros_like(truncated)
-            if terminated.any():
+            if traj_terminated.any():
                 if isinstance(seq_length, int):
                     truncated.view(num_slices, -1)[traj_terminated] = 1
                 else: