Add support for IterableDatasets everywhere

[ethanwharris]

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you write any new necessary tests?

What does this PR do?

Fixes #948

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[Borda]

I would rather create complet Dataloader so it is easier to undestand... what about?

class CustomInfDataloader:

    def __init__(self, dataset, batch_size, shuffle):
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle

    def __iter__(self):
        idxs = []
        while True:
            if len(idxs) < self.batch_size:
                idxs = range(len(self.dataset))
                if self.shuffle:
                    np.random.shuffle(idxs)
            batch = [self.dataset[idx] for idx in idxs[:self.batch_size]]
            yield batch
            idxs = idxs[len(batch):]

[ethanwharris]

torch.DataLoader does quite a bit more than this (e.g. collate functions, samplers, etc.) so it is probably better to wrap it rather than rewrite it - also we don't really have access to the dataset when this is created, only the dataloader

[williamFalcon]

@Borda we generally want to avoid duplicating torch functionality. Otherwise the project scope will blow up quickly,

[Borda]

I do agree, I just found this construction quite difficult to follow...

[Borda]

hey there, we have added GPU CI test, so could we kindly ask to rebase/merge master which will trigger these tests so we do not need to test it manually... Thx for your understanding 🤖

[ethanwharris]

@Borda Done :)