Hello, As mentioned here torchvision ops is currently not supported in torchscript. I wanted to know is there a timeline on when will torchscript support these functions.
The documentation might be slightly out of date if you compile both PyTorch and torchvision yourself or take the nightlies, because it seems to work for me:
In [1]: import torch, torchvision
In [2]: @torch.jit.script
...: def fn(a, b):
...: return torchvision.ops.nms(a, b, 0.3)
...:
In [3]: fn.graph
Out[3]:
graph(%a.1 : Tensor,
%b.1 : Tensor):
%5 : Function = prim::Constant[name="nms"]()
%4 : float = prim::Constant[value=0.29999999999999999]() # <ipython-input-2-cfe32fda8960>:3:37
%6 : Tensor = prim::CallFunction(%5, %a.1, %b.1, %4) # <ipython-input-2-cfe32fda8960>:3:11
return (%6)
In [4]: fn(torch.randn(10, 4), torch.randn(10))
Out[4]: tensor([2, 8, 6, 7, 0, 9, 4, 5, 3, 1])
Best regards
Thomas
@tom Thanks. It is working for me now in torchscript as well. But, the NMS now fails in the TRTserver.
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.
Well, you might need to load the custom ops for that. I know nothing about TRTserver.