[FIXED] return next(self._sampler_iter) # may raise StopIteration

Issue

Instead of using enumerate(data loader) for some reason, I am creating iterator for the data loader. In the while loop shown below, it gives me StopIteration error.

Minimalistic code that depicts the cause:

loader = DataLoader(dataset, batch_size=args.batch_size)
dataloader_iter = iter(loader)
while(dataloader_iter):
    X, y = next(dataloader_iter)
    ...

What would be the correct condition (to specify within the while loop) to check if the iterator is empty?

Solution

In Python it's standard in a lot of cases to use exceptions for control flow.

Just wrap it in a try-except:

loader = DataLoader(dataset, batch_size=args.batch_size)
dataloader_iter = iter(loader)
try:
    while True:
        x, y = next(dataloader_iter)
        ...
except StopIteration:
    pass

If you want to catch some other errors inside the while loop, you can move the try-except inside, but you must remember to break out of the loop when you hit a StopIteration:

loader = DataLoader(dataset, batch_size=args.batch_size)
dataloader_iter = iter(loader)
while True:
    try:
        x, y = next(dataloader_iter)
        ...
    except SomeOtherException:
        ...
    except StopIteration:
        break

Answered By - flakes