Hi,
Yes, indeed the answer from my thread can pose a problem of uneven sampling. However, I found alternatives on setting manual step within epoch while still make dataloader goes over all data in dataset.
You could make another class that acts as wrapper over dataloader, in which you manually set how much step needed for given dataloader and number of steps
class LoaderWrapper:
    def __init__(self, dataloader, n_step):
        self.step = n_step
        self.idx = 0
        self.iter_loader = iter(loader)
    
    def __iter__(self):
        return self

    def __len__(self):
        return self.step

    def __next__(self):
        # if reached number of steps desired, stop
        if self.idx == self.step:
            self.idx = 0
            raise StopIteration
        else:
            self.idx += 1
        # while True
        try:
            return next(self.iter_loader)
        except StopIteration:
            # reinstate iter_loader, then continue
            self.iter_loader = iter(self.loader)
            return next(self.iter_loader)

Using toy sample, here’s what I got for using standard DataLoader

While this is my result using wrapper above

This might be less elegant that what you might expect, but hopefully this can solve your problem
Cheers