It’s not expected and you are correct finding this issue! It looks like I’ve used the wrong code, as I first wanted to define train_model like a closure, i.e. basically training on one batch, but apparently haven’t cleaned up my code. Sorry about that.


Your new code looks alright. I would just re-create the optimizer after resetting the model.


output will contain the logits for each single test sample. So if you would like to print the probabilities for each class, you could use print('Class probabilities: {}'.format(F.softmax(output, 1))).


Could you try to fix point 2 (re-creating the optimizer)? If the code still gets stuck, I’ll try to reproduce it.