[BugFix] Extend with a list of tensordicts (#2032)

pytorch · Mar 21, 2024 · a69c667 · a69c667
1 parent 660d827
commit a69c667
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/test/test_rb.py b/test/test_rb.py
@@ -775,6 +775,12 @@ class TC:
         if data == "tc":
             assert storage._storage.text == storage_recover._storage.text
 
+    def test_add_list_of_tds(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(100))
+        rb.extend([TensorDict({"a": torch.randn(2, 3)}, [2])])
+        assert len(rb) == 1
+        assert rb[:].shape == torch.Size([1, 2])
+
 
 @pytest.mark.parametrize("max_size", [1000])
 @pytest.mark.parametrize("shape", [[3, 4]])

diff --git a/torchrl/data/replay_buffers/storages.py b/torchrl/data/replay_buffers/storages.py
@@ -1527,6 +1527,8 @@ def save_tensor(tensor_path: str, tensor: torch.Tensor):
 
 
 def _flip_list(data):
+    if all(is_tensor_collection(_data) for _data in data):
+        return torch.stack(data)
     flat_data, flat_specs = zip(*[tree_flatten(item) for item in data])
     flat_data = zip(*flat_data)
     stacks = [torch.stack(item) for item in flat_data]