How to customize the module `state_dict` & `load_state_dict`

I don’t know if there is a better approach. While internally state_dict hooks are used, they can change at any time and could thus easily break your code if you depend on them.