Python ScriptModule: Could not cast attribute 'arch_opt' to type __torch__.modules.ArchitectureOptions

[yuanyaoshen]

I am trying to convert a pytorch module to ScriptModule. What I was trying to do is adding

nl_denoiser_script = torch.jit.script(nl_denoiser)
print(nl_denoiser_script.code)

after nl_denoiser.patch_denoise_net.load_state_dict(model_state0['state_dict']) in blind_denoise_bw_small.py. However, I ran into issue like:

RuntimeError: Could not cast attribute 'arch_opt' to type torch.modules.ArchitectureOptions (of Python compilation unit at: 0x518f410): Could not cast attribute 'rgb' to type Tensor: Unable to cast Python instance to C++ type (compile in debug mode for details)

I tried adding @torch.jit.script in front of class ArchitectureOptions but the same error occurred. Is there solution for this? Thank you!