torch.optim.Optimizer.register_state_dict_post_hook¶
- Optimizer.register_state_dict_post_hook(hook, prepend=False)[source][source]¶
Register a state dict post-hook which will be called after
state_dict()
is called.It should have the following signature:
hook(optimizer, state_dict) -> state_dict or None
The hook will be called with arguments
self
andstate_dict
after generating astate_dict
onself
. The hook may modify the state_dict inplace or optionally return a new one. The registered hook can be used to perform post-processing on thestate_dict
before it is returned.- Parameters
hook (Callable) – The user defined hook to be registered.
prepend (bool) – If True, the provided post
hook
will be fired before all the already registered post-hooks onstate_dict
. Otherwise, the providedhook
will be fired after all the already registered post-hooks. (default: False)
- Returns
a handle that can be used to remove the added hook by calling
handle.remove()
- Return type
torch.utils.hooks.RemoveableHandle