I think the behavior might be related to scripting the module.
If you are using @torch.jit.script_method, you should derive the module from torch.jit.ScriptModule.
Could you update PyTorch to the nightly builds (have a look at the instructions 9) and run the code again?