Solved by albanD in post #2 
              
              
                Hi, 
For me, “epoch” has a slightly different meaning: it means how many times you go through the whole dataset. With this, I would do: 
n_smaples_seen = 0
for epoch in range(num_epochs):
    for sample in dl:
        # Not needed if your code handle any batch size
        if sample.size(0) != bs:
 …