Solved by tom in post #2 
              
              
                The easiest way would be to just iterate over subepochs and break after the number of iterations you want. 
for e in range(num_epochs):
  it = 0
  while it < num_it:
    for batch in data_loader:
      it += 1
      if it >= num_it:
        break
      ....

It’s a tad verbose, but in the end at 5 a…