Torchvision.ops.nms in torchscript

jit
2

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