I believe that the problem resides in the fact that I am removing param_groups as following:
del optimizer.param_groups[1]
This is removing the second param group from the optimizer but it seems that it’s associated parameters in optimizer.state reside leading to an error when matching param_mappings with optimizer.state.
If that is the case then what is the right way to remove param_groups? (no documentation regarding this in PyTorch optimizer documentation)