Another approach would be to set an arbitrary high number as length and if you encounter the index to be to high you can choose of two options:


You can modify the index to be in your range (e.g. By using the % operator. This would result in a fixed epoch length but probably iterating much more often over some parts of the dataset (should not matter if you shuffle anyway)


You can raise the StopIteration yourself which should be caught inside the Dataloader/automatically inside the loop. This is what I did in early releases of pytorch to write own imagefolder-like datasets (if it is not caught you need to catch it yourself)


Note that for the second approach might cause loggers or other modules which are build upon a fixed number of batches to raise errors.