The issue is that the stored probs in lists are invalid after the first optimizer.step() and parameter update.
I.e. the prob tensors were creates with a model using the initial parameter set P0.
After the first loss.backward() and optimizer.step() operation these parameters will be updated to P1.
In the next loss.backward() and optimizer.step() iteration you are trying to calculate the gradients of the “old” loss w.r.t the new parameters, which is wrong since these already parameters didn’t calculate the loss.
Older PyTorch versions didn’t implement this check and would thus yield wrong results.
The proper way would be to recalculate the prob after each update or to calculate all gradients before calling optimizer.step().