Solved by ptrblck in post #3 
              
              
                The last batch given by your DataLoader might be smaller, if the length of your Dataset is not divisible by the batch size without a remainder. 
Try to get the current batch size inside the training loop: 
for data, target in testloader:
    batch_size = data.size(0)
    ...
    for i in range(batch…