Pytorch: use an IterableDataset to create an infinite stream of batched data

infinite_dataset.py

from torch.utils.data import IterableDataset, DataLoader

# The dataset has no length
class TestDataset(IterableDataset):

    def __iter__(self): # __iter__() can be a generator
        while True:
            yield 5

# The loader is called as normal
loader = DataLoader(TestDataset(), batch_size=3)

for b in loader:
    print(b)

# Result: repeats [5, 5, 5] forever