Even though you retrieve the previous model, I think you are still updating the internal state of your optimizer when calling optim.step(). Adam has running estimates of different moments which are updated at each call to step(). The same is true for e.g. SGD with momentum.
Therefore, the next call to optim.step() will not give the same results (0.1989, 0.1696) even if the loss and everything else is the same. You could try to also store the optimizer’s state_dict and reload that as part of the “back-stepping”.