amend

pytorch · Feb 7, 2024 · 91941ea · 91941ea
1 parent fe960fa
commit 91941ea
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchrl/data/replay_buffers/storages.py b/torchrl/data/replay_buffers/storages.py
@@ -911,7 +911,8 @@ def _init(self, data: Union[TensorDictBase, torch.Tensor]) -> None:
             out = out.expand(self.max_size, *data.shape)
             out = out.memmap_like(prefix=self.scratch_dir)
             if self.names is not None:
-                out.names = self.names
+                names = self.names + [None] * (out.batch_dims - len(self.names))
+                out.refine_names(*names)
             for key, tensor in sorted(
                 out.items(include_nested=True, leaves_only=True), key=str
             ):