[BUG] DataLoader with LazyStackedTensorDict of different sizes

[Modexus]

Describe the bug

Not sure if I am missing something but I am trying to load a batch of lazy_stacked TensorDicts that have variable size tensors.
The DataLoader throws an error however that it cannot stack the tensors.

To Reproduce

import tensordict
import torch

tensors = [{"x": torch.rand((i,))} for i in range(10)]
tensordicts_stacked = tensordict.lazy_stack(
    [tensordict.TensorDict.from_dict(x) for x in tensors]
)

dl = DataLoader(tensordicts_stacked, batch_size=4, collate_fn=lambda x: x)
next(iter(dl))

Expected behavior

Not sure if this is supposed to work but I would have expected it to return a batch of LazyStackedTensorDicts.

Reason and Possible fixes

I think it's because __getitems__ actually points to <bound method TensorDictBase.__getitem__ of LazyStackedTensorDict.
If I add __getitems__ = __getitem__ after the __getitem__ in LazyStackedTensorDict (the same as in the Base) it returns a LazyStackedTensorDict without an error.

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)