[BUG] tensordict.pad_sequence with a tensorclass returns TensorDict

[egaznep]

Describe the bug

The function tensordict.pad_sequence is great for collating individual datapoints of varying length. And tensorclasses are good for typing hints. However, these two do not work together nicely because pad_sequence always returns a TensorDict and never the given class.

To Reproduce

import torch
import tensordict
@tensordict.tensorclass
class DataPoint:
    a: torch.Tensor
    b: torch.Tensor

s1 = DataPoint(a=torch.ones(10, 1), b=torch.ones(10, 1), batch_size=[])
s2 = DataPoint(a=torch.ones(20, 1), b=torch.ones(30, 1), batch_size=[])
tensordict.pad_sequence([s1, s2], batch_first=True)

Output: 
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([2, 20, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([2, 30, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)

Expected behavior

Ideally, return should be of the type DataPoint.

Reason and Possible fixes

In pad_sequence,

    if out is None:
        out = TensorDict(
            {}, batch_size=torch.Size(shape), device=device, _run_checks=False
        )

constructs the return object. Either a modification here could solve it, or right before returning the TensorDict, we can use the tensorclass._from_tensordict constructor (in our case, DataPoint._from_tensordict) to comply with the original data type.

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)

[vmoens]

Yeah we should call empty there! Good catch

[vmoens]

Edit: On nightly I get a Datapoint object, which version are you using?

[egaznep]

Edit: On nightly I get a Datapoint object, which version are you using?

Oh really? I had
tensordict 0.1.2+6de2db2.

as available on conda-forge https://anaconda.org/conda-forge/tensordict

There's 0.4.0 out?!

[vmoens]

Yeaaah not very good at upgrading conda :/ Let me do that!

[egaznep]

Thank you! I guess this issue could be closed :)