PyTorch DataLoader improvements for Iterable Dataset about pytorch HOT 4 OPEN

Here is an example resulting of this problem where the IterableDataset is un-aware of the batch_size and num_workers: Lightning-AI/pytorch-lightning#19624

from pytorch.

@tchaton thanks for flagging this! Definitely understand and agree that IterableDatasets are difficult to get correct. Just to check my understanding: I think the proposal you're making is to pass information from DataLoader to dataset through a method call, but the user would still be required to set the behaviour up to eg end iteration early.

To be honest I'm not totally convinced how much of an improvement this is. Re: the linked Lightning issue, that does sound like a pain but I'm not sure how this proposal would help in this situation.

from pytorch.

Hey @andrewkho Yes, there is a lot of things hard to get it right.

As an example, here: https://github.com/Lightning-AI/litdata/blob/f23ba626c48d62a8397c2b7ae4bacfd345df2c86/src/litdata/streaming/dataloader.py#L620.

I had to override the length of the DataLoader to ensure the length was correct

from pytorch.

Here is another example of this issue: Lightning-AI/litdata#147

from pytorch.