Below is the source code of _register_state_dict_hook I quoted from official Pytorch nn.Module
def _register_state_dict_hook(self, hook):
r"""These hooks will be called with arguments: `self`, `state_dict`,
`prefix`, `local_metadata`, after the `state_dict` of `self` is set.
Note that only parameters and buffers of `self` or its children are
guaranteed to exist in `state_dict`. The hooks may modify `state_dict`
inplace or return a new one.
"""
handle = hooks.RemovableHandle(self._state_dict_hooks)
self._state_dict_hooks[handle.id] = hook
return handle
The comment says that “These hooks will be called … after the “state_dict” of self is set”.
While I’m not sure when exactly the “state_dict” of a nn.Module is set. Is it set after initialization? Or every time after the weight get updated?