Torchvision.ops.nms in torchscript

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                                                                                                                                                                                                   
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


@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.