It’s on the ds_loss backward (you can see it in the second snippet inside the inner and after the  # DIFF comment).

Because the acc_loss will backprop in the same graph as the ds_loss so I think the retain_graph is required here.

I thought about it and I had also tried to change ds_loss.backward(retain_graph=True) to criterion(fmodel(ds_imgs, ds_trgs)).backward() to create a separate path on the graph which does not interfere with the following backpropagation on acc_loss. Unfortunately, that didn’t help much as the problem occur earlier when I call for the second time ds_loss.backward(). Since every iteration optimizes a different learning rate I don’t understand why should I retain the graph there.

For the slowdown, I guess it depends how many times you unroll your loop as you differentiate through all these iterations.

That’s fair but for the same number of inner iterations, the second snippet takes a lot more. I don’t think it should since I’m learning literally just one new parameter and it doesn’t have to backprop through several steps (just through one since I update it after I evaluate the model on the training set).
I think that whatever the bug is the two “symptoms” are correlated i.e. it’s taking so much cause it’s doing something which requires retain_graph=True while what I’m actually trying to do shouldn’t need it and should do a much more simple (and cost efficient) update