[BugFix] Fix Atari DQN ensembling (#1981)

Author: Vincent Moens

torchrl/data/datasets/atari_dqn.py

@@ -745,6 +745,9 @@ def get(self, index):
         if isinstance(index, torch.Tensor):
             if index.ndim <= 1:
                 return self._read_from_splits(index)
+            elif index.shape[1] == 1:
+                index = index.squeeze(1)
+                return self.get(index)
             else:
                 raise RuntimeError("Only 1d tensors are accepted")
             # with ThreadPoolExecutor(16) as pool:

torchrl/data/replay_buffers/samplers.py

@@ -1565,7 +1565,7 @@ def sample(self, storage, batch_size):
                 ]
             )
         samples = [
-            sample if isinstance(sample, torch.Tensor) else torch.tensor(sample)
+            sample if isinstance(sample, torch.Tensor) else torch.stack(sample, -1)
             for sample in samples
         ]
         if all(samples[0].shape == sample.shape for sample in samples[1:]):