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: