fix lazy stack dataloader batching

Author: gtnv

tensordict/_lazy.py

@@ -2696,6 +2696,8 @@ def recompose(converted_idx, stack_dim=new_stack_dim):
                     result.lock_()
                 return result
 
+    __getitems__ = __getitem__
+
     def __eq__(self, other):
         return self._dispatch_comparison(other, "__eq__", "__eq__", default=False)
 

test/tensordict/test_lazy.py

@@ -275,6 +275,17 @@ def test_to_mds(self, tmpdir):
         test_td = LazyStackedTensorDict(*batches)
         assert_allclose_td(td, test_td)
 
+    def test_dataloader_ragged_lazy_stack(self):
+        dataset = lazy_stack([TensorDict(x=torch.zeros(i)) for i in range(10)])
+        batches = list(
+            torch.utils.data.DataLoader(  # noqa: TOR401
+                dataset, batch_size=4, collate_fn=lambda batch: batch
+            )
+        )
+        assert [len(batch) for batch in batches] == [4, 4, 2]
+        assert all(isinstance(batch, LazyStackedTensorDict) for batch in batches)
+        assert [len(td["x"]) for batch in batches for td in batch] == list(range(10))
+
     def test_all_keys(self):
         td = TensorDict({"a": torch.zeros(1)}, [])
         td2 = TensorDict({"a": torch.zeros(2)}, [])