IngoMarquart:

Using a custom batch sampler and a map style dataset, or using an iterative dataset, where the iteration is determined by a custom batch size.


Yes, you are right that there are two ways:

use IterableDataset: sampling is in dataset code, which is run in worker process.
use batch_size=None and sampler=custom_sampler_that_is_really_a_batch_sampler: sampling is done by a separate class, run in main process.

In general, IterableDataset gives you finer control, because you can achieve things like data-dependent batch size / ordering. And depending on your code, having sampling logic and fetching logic in the same class may be cleaner. But since sampling is done in workers, you likely need to configure each worker to make sure they don’t load duplicate batches, or load in correlated ordering. To me, the second approach makes more sense if your sampling and fetching logic are relatively isolated, and is usually less tricky to implement.
Re docs: The docs could use some improvements… It was largely my fault. If I find some free time I would update it. Sorry about it.