If you want to create an ‘infinite generator’ simply place it in a while loop:
def cycle(iterable):
while True:
for x in iterable:
yield x
However there is a more elegant solution with a DataLoader. Simply enumerate over the dataloader when training for one epoch:
data_loader = torch.utils.data.DataLoader(...)
def train(epoch):
for batch_idx, (data, target) in enumerate(data_loader):
# Your code for training one epoch.
# Now perform training for however many epochs you want.
for epoch in range(num_epochs):
train(epoch)
Internally every time you call enumerate(data_loader), it will call the object’s iter method which will create and return a new _DataLoaderIter object, which internally calls iter() on your batch sampler here.