[BugFix] Fix displaying of tensor sizes in buffers

ghstack-source-id: 511609a169996b680dccbd272e3d5b5710618558
Pull Request resolved: #2456

torchrl/data/replay_buffers/storages.py

@@ -24,6 +24,7 @@
     TensorDict,
     TensorDictBase,
 )
+from tensordict.base import _NESTED_TENSORS_AS_LISTS
 from tensordict.memmap import MemoryMappedTensor
 from torch import multiprocessing as mp
 from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
@@ -1120,7 +1121,12 @@ def max_size_along_dim0(data_shape):
             out = out.memmap_like(prefix=self.scratch_dir, existsok=self.existsok)
             if torchrl_logger.isEnabledFor(logging.DEBUG):
                 for key, tensor in sorted(
-                    out.items(include_nested=True, leaves_only=True), key=str
+                    out.items(
+                        include_nested=True,
+                        leaves_only=True,
+                        is_leaf=_NESTED_TENSORS_AS_LISTS,
+                    ),
+                    key=str,
                 ):
                     try:
                         filesize = os.path.getsize(tensor.filename) / 1024 / 1024