I try to use torchvision.ops.nms in my model. and i have used the following to save my model using torch.jit.trace.
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("model.pt")
Everything works well. But when load model.pt on mobile,came across error.
E/AndroidRuntime: Caused by: com.facebook.jni.CppException:
Unknown builtin op: torchvision::nms.
Could not find any similar ops to torchvision::nms. This op may not exist or may not be currently supported in TorchScript.
:
.................
/usr/lib/python3.6/runpy.py(85): _run_code
/usr/lib/python3.6/runpy.py(193): _run_module_as_main
Serialized File "code/__torch__/resnet.py", line 68
_39 = torch.slice(torch.select(scores0, 0, 0), 0, 0, 9223372036854775807, 1)
_40 = torch.slice(_38, 1, 0, 9223372036854775807, 1)
anchors_nms_idx = ops.torchvision.nms(_40, torch.select(_39, 1, 0), 0.5)
~~~~~~~~~~~~~~~~~~~ <--- HERE
_41 = torch.slice(torch.select(classification0, 0, 0), 1, 0, 9223372036854775807, 1)
anchors_nms_idx0 = torch.to(anchors_nms_idx, dtype=4, layout=0, device=torch.device("cpu"), pin_memory=False, non_blocking=False, copy=False, memory_format=None)
Can anyone please help me? Thank you