amend

pytorch · Mar 4, 2024 · ed98a4f · ed98a4f
1 parent ef7d366
commit ed98a4f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrl/data/datasets/atari_dqn.py b/torchrl/data/datasets/atari_dqn.py
@@ -761,8 +761,8 @@ def _read_from_splits(self, item: int | torch.Tensor):
         split = (item < self.frames_per_split[1:, 1].unsqueeze(1)) & (
             item >= self.frames_per_split[:-1, 1].unsqueeze(1)
         )
-        # split_tmp, idx = split.squeeze().nonzero().unbind(-1)
-        split_tmp, idx = split.nonzero().unbind(-1)
+        split_tmp, idx = split.squeeze(0).nonzero().unbind(-1)
+        # split_tmp, idx = split.nonzero().unbind(-1)
         split = torch.zeros_like(split_tmp)
         split[idx] = split_tmp
         split = self.frames_per_split[split + 1, 0]