state_dict_hook is an internal implementation that allows some modules to edit their state_dict before they get returned. torch.jit doesn’t generally support any hooks, including state_dict hooks, so it is possible that Apex is installing these hooks on a jit module and then it is failing. However, I do not know the details of how Apex works, so I am not sure that is precisely what is happening.