[BUG] dataloader only return the first element of the sequence

[wjfwzzc]

🐛🐛 Bug Report

⚗️ Current Behavior

When I build a dataset with a sequence of images (like frames), everything works fine. But if I tranfer it to a pytorch dataloader, the loader can only return the first element of the sequence.

Input Code "Creating Datasets with Sequences" official colab example

Append following code to the end:

print(ds[0]["frames"].shape)
print(ds[1]["frames"].shape)
dataloader = ds.pytorch()
data_iter = iter(dataloader)
print(next(data_iter)["frames"].shape)
print(next(data_iter)["frames"].shape)

It returns:

(600, 1080, 1920, 3)
(1050, 1080, 1920, 3)
torch.Size([1, 1080, 1920, 3])
torch.Size([1, 1080, 1920, 3])

But I hope It returns something like:

(600, 1080, 1920, 3)
(1050, 1080, 1920, 3)
torch.Size([1, 600, 1080, 1920, 3])
torch.Size([1, 1050, 1080, 1920, 3])