Can`t load JIT model from .trcd file

[paxhx]

I saved 3D DensNet from monai:

traced = torch.jit.trace(model, torch.rand(1, 3, 160, 160, 160).cuda())
torch.jit.save(traced, f'monai3_160_baseline.trcd')

Now I am trying to load it from trcd file with the following function:

torch.jit.load(x, map_location="cuda:0")

And I get the following error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-28-60944cabc718> in <module>()
----> 1 model = load_jit_model("./model_training/monai3_160_baseline.trcd").cuda()

<ipython-input-22-2020de3a15ce> in load_jit_model(x, cuda_id)
     14 
     15 def load_jit_model(x, cuda_id=0):
---> 16     return torch.jit.load(x, map_location=f'cuda:{cuda_id}' if torch.cuda.is_available() else 'cpu')

/home/yeliseev/anaconda3/envs/rsna/lib/python3.7/site-packages/torch/jit/__init__.py in load(f, map_location, _extra_files)
    273     cu = torch._C.CompilationUnit()
    274     if isinstance(f, str) or isinstance(f, pathlib.Path):
--> 275         cpp_module = torch._C.import_ir_module(cu, f, map_location, _extra_files)
    276     else:
    277         cpp_module = torch._C.import_ir_module_from_buffer(cu, f.read(), map_location, _extra_files)

RuntimeError: expected ident but found 'class' here:
Serialized   File "code/__torch__/torch/nn/modules/container/___torch_mangle_423.py", line 8
  norm : __torch__.torch.nn.modules.pooling.AdaptiveAvgPool3d
  flatten : __torch__.torch.nn.modules.flatten.Flatten
  class : __torch__.torch.nn.modules.linear.Linear
  ~~~~~ <--- HERE
  def forward(self: __torch__.torch.nn.modules.container.___torch_mangle_423.Sequential,
    argument_1: Tensor) -> Tensor:

[paxhx]

I have the following versions of torch and monai:
torch==1.6.0
monai==0.3.0

[Nic-Ma]

Hi @YuraYelisieiev ,

Thanks for your interest here.
May I know what you get if change to save in .pt format(monai3_160_baseline.pt)?

Thanks.

[paxhx]

@Nic-Ma Its the same for .pt file

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-11-cdc5fc4ecacb> in <module>()
----> 1 model = load_jit_model("./model_training/monai3_160_baseline.pt").cuda()

<ipython-input-10-2020de3a15ce> in load_jit_model(x, cuda_id)
     14 
     15 def load_jit_model(x, cuda_id=0):
---> 16     return torch.jit.load(x, map_location=f'cuda:{cuda_id}' if torch.cuda.is_available() else 'cpu')

/home/yeliseev/anaconda3/envs/rsna/lib/python3.7/site-packages/torch/jit/__init__.py in load(f, map_location, _extra_files)
    273     cu = torch._C.CompilationUnit()
    274     if isinstance(f, str) or isinstance(f, pathlib.Path):
--> 275         cpp_module = torch._C.import_ir_module(cu, f, map_location, _extra_files)
    276     else:
    277         cpp_module = torch._C.import_ir_module_from_buffer(cu, f.read(), map_location, _extra_files)

RuntimeError: expected ident but found 'class' here:
Serialized   File "code/__torch__/torch/nn/modules/container/___torch_mangle_423.py", line 8
  norm : __torch__.torch.nn.modules.pooling.AdaptiveAvgPool3d
  flatten : __torch__.torch.nn.modules.flatten.Flatten
  class : __torch__.torch.nn.modules.linear.Linear
  ~~~~~ <--- HERE
  def forward(self: __torch__.torch.nn.modules.container.___torch_mangle_423.Sequential,
    argument_1: Tensor) -> Tensor:

[Nic-Ma]

Hi @YuraYelisieiev ,

OK, I see.
Thanks very much for your bug report.
I submitted PR: #1114 to fix this issue.
Could you please help review and verify it?

Thanks.

[paxhx]

Yes, no problem. Thanks for the fix.