Custom dataset with unknown length

I have a custom mux dataset that hold k different datasets.

I train my model with iteration rather than epochs so I want to make sure the mux dataset is infinite dataset object since the dataloader sometimes get stuck when I recreate the iterator using the following:

def endless_iterator(iterable: Iterator[any]):
    while True:
        for elem in iterable:
            yield elem

Now a valid solution is to set the len to return a very large number but I was wondering if there is a way to just make it work endlessly with next(data_loader) without the need of my endless_iterator.

In case you are using num_workers>0 you could use persistent_workers=True, which would keep the workers alive without tearing them down at the end of an epoch.