amend

torchrl/data/replay_buffers/storages.py

@@ -283,11 +283,6 @@ def _len(self, value):
             _len_value = self._len_value = mp.Value("i", 0)
         _len_value.value = value
 
-    def __getitem__(self, item):
-        if self._len < self.max_size:
-            return self._storage[: self._len][item]
-        return super().__getitem__(item)
-
     def __getstate__(self):
         state = copy(self.__dict__)
         if get_spawning_popen() is None:
@@ -453,11 +448,15 @@ def set(  # noqa: F811
         self._storage[cursor] = data
 
     def get(self, index: Union[int, Sequence[int], slice]) -> Any:
+        if self._len < self.max_size:
+            storage = self._storage[: self._len]
+        else:
+            storage = self._storage
         if not self.initialized:
             raise RuntimeError(
                 "Cannot get an item from an unitialized LazyMemmapStorage"
             )
-        out = self._storage[index]
+        out = storage[index]
         if is_tensor_collection(out):
             out = _reset_batch_size(out)
             return out.unlock_()