model.train()
for inputs, labels in train_loader:
The model.train() needs to go there. If you put it outside as in your snippet, the model will only be in training mode for the first epoch. All subsequent epochs will be in evaluation mode.
1 Like