_register_state_dict_hook is not supported on ScriptModules

jit
1

I’m quite new to torch.jit and Apex. I’m getting this error. What does it mean ?

RuntimeError: _register_state_dict_hook is not supported on ScriptModules

Here’s my code

model, optimizer = amp.initialize(model, optimizer,opt_level='O2',
                                          keep_batchnorm_fp32=False,
                                          loss_scale=config.TRAIN.FP16_LOSS_SCALE,
                                          min_loss_scale=128.0)
2

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.

3

Scripting is currently not supported in apex, but tracing should work.
We are currently working on upstreaming mixed-precision training, so that you won’t need to build apex soon. :wink:

1 Like