I reviewed this once more, and it turns out setting the parameter back to None is enough to unregister it again, so the issue I raised earlier is invalid.

I rebased and made a small edit, where I also fixed the styling of the long line and used the normal way to assign attributes, since this is equivalent to setattr.

@chr5tphr btw, another idea is to mark new parameters so that they can be deleted on remove

For example:

#...
for module in modules:
    if module.bias is None:
        module.bias = torch.nn.Parameter(torch.zeros(1, device=module.weight.device, dtype=module.weight.dtype))
        module.bias.is_temporary = True

Thanks for the reply!
I verified that setting an attribute that is a parameter to None does unregister it, so we do not need to explicitly delete it, as we store the parameter regardless of whether it was a torch.nn.Parameter, or None.