Hi,
The main problem you encounter here is that the nn.Module and torch.optim modules are not built to do what you want. In particular, the parameters of the modules and the optimizers are built such that weight updates are not differentiated through.
The two issues you want to fix here are:

You don’t want to use builtin optimizers as they perform their updates in a non-differentiable way.
You don’t want to use vanilla modules because their weights are nn.Parameter for which you never track history (they are always leaf Tensors).

If you only want to perform testing for research purposes, I would advise

Create regular pytorch modules.
Create the new base parameters as mod.true_weight = nn.Parameter(mod.weight)

Delete all the old parameters by doing del mod.weight manually.
Have a way to manually set the .weight to the base parameters (that will use only for your first iteration)

For every param, do mod.weight = mod.true_weight.view_as(mod.true_weight) (do an operation to make sure you don’t get a new Tensor (that share memory)).


Have a custom optimizers that will do in a differentiable way the weight update based on the true_weight’s gradients: mod.weight = mod.weight + lr * mod.true_weight.grad.
Use the regular mod(inp) to do the forward calls.

That way, you will be able to manually do these different iterations.