`__getitem__()` does not support returning tensors with different sizes

2

You could use a custom collate_fn as described here which would work for variable shapes.