I’m not sure to understand then.
Do you want the gradient computation to take into account the first forward/backward, the optimizer step and the second forward?
Or you just want to use gradients computes with W2 and apply them to W1?