[BugFix] Fix lazy-stack in RBs

ghstack-source-id: 38399ee
Pull Request resolved: #2880

Author: vmoens

torchrl/data/replay_buffers/storages.py

@@ -431,7 +431,7 @@ def get(self, index: int | Sequence[int] | slice) -> Any:
             stack_dim = self.stack_dim
             if stack_dim < 0:
                 stack_dim = out[0].ndim + 1 + stack_dim
-            out = lazy_stack(list(out), stack_dim=stack_dim)
+            out = lazy_stack(list(out), stack_dim)
             return out
         return out
 

torchrl/data/replay_buffers/utils.py

@@ -17,7 +17,7 @@
 import numpy as np
 import torch
 from tensordict import (
-    LazyStackedTensorDict,
+    lazy_stack,
     MemoryMappedTensor,
     NonTensorData,
     TensorDict,
@@ -482,7 +482,7 @@ def __call__(self, data: TensorDictBase, out: TensorDictBase = None):
                         out._get_sub_tensordict((slice(None),) * i + (j,))
                         for j in range(out.shape[i])
                     ]
-                    out = LazyStackedTensorDict(*out_list, stack_dim=i)
+                    out = lazy_stack(out_list, i)
 
             # Create a function that reads slices of the input data
             with out.flatten(1, -1) if out.ndim > 2 else contextlib.nullcontext(