So this is going to be a bit more tricky.
The first backward with W1 will need to set create_graph=True.
But then you won’t be able to use optimizer to do the step as they operate inplace and without registering their operation to the autograd.
So you either need to reimplement the optimizer to do this operation in a differentiable manner. And assign these new values without deleting the old ones and keeping gradients flowing (which can be tricky).
Or check libraries such as higher 11 that already do similar things