register_module_load_state_dict_pre_hook

→RemovableHandle

register_module_load_state_dict_pre_hook(hook: Callable[..., object])

source

Invoked before a module attempts to copy values from a state-dict into its own parameters / buffers. The hook may mutate state_dict in place to rename keys, transform tensors (e.g. cast dtype, adjust shape after architectural changes), or pre-populate missing_keys/unexpected_keys to influence strict=True behaviour.

Parameters

hookCallable

Signature::

hook(module, state_dict, prefix, local_metadata, strict,
     missing_keys, unexpected_keys, error_msgs) -> None

All list / dict arguments may be mutated in place.

Returns

RemovableHandle

Handle for later deregistration.

Notes

Common application: backward-compatible checkpoint loading after refactors — e.g. rename "layer.weight" to "linear.weight" to accommodate a renamed submodule without forcing users to regenerate checkpoints.

Examples

>>> from lucid.nn.hooks import register_module_load_state_dict_pre_hook
>>> def rename_keys(mod, sd, prefix, *args):
...     # Rename legacy keys on the fly.
...     for k in list(sd):
...         if k.startswith(prefix + 'fc.'):
...             sd[k.replace('fc.', 'linear.', 1)] = sd.pop(k)
>>> h = register_module_load_state_dict_pre_hook(rename_keys)
>>> # model.load_state_dict(legacy_checkpoint) now succeeds