Problems with jit load 3D monai model

I saved 3D DensNet from monai:

traced = torch.jit.trace(model, torch.rand(1, 3, 160, 160, 160).cuda()), 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)
     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/ 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,, map_location, _extra_files)

RuntimeError: expected ident but found 'class' here:
Serialized   File "code/__torch__/torch/nn/modules/container/", 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:

Which PyTorch an MONAI version are you using?


Thanks for your reply.