amend

pytorch · Mar 4, 2024 · 4860c01 · 4860c01
1 parent 73ce0d7
commit 4860c01
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/data/datasets/atari_dqn.py b/torchrl/data/datasets/atari_dqn.py
@@ -673,7 +673,7 @@ def preprocess(
             first_item = self[0]
             mmap = fn(first_item)
             mmap = mmap.expand(len(self), *first_item.shape)
-            mmap["_indices"] = torch.arange(mmap.shape[0])
+            mmap["_indices"] = torch.arange(mmap.shape[0]).reshape(mmap.shape)
             mmap.memmap_like(tmpdir, num_threads=32)
 
             def func(mmap):