Repository: pytorch
Issue: Unknown builtin op: torchvision::nms when loading scripted FasterRCNN
Labels: oncall: jit

## 🐛 Bug

Scripted FasterRCNN saved to a .pt-file cannot be loaded via torch.script.load due to 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.

## To Reproduce

Steps to reproduce the behavior:

model.eval()
model = torch.jit.script(model)
torch.jit.save(model, PATH_MODEL + 'model.pt')
model = torch.jit.load('./models/model.pt') # error occurs

## Expected behavior

Model should load without error

## Environment

 - PyTorch Version (e.g., 1.0): 1.7.0+cu110
 - OS (e.g., Linux): Windows 10
 - How you installed PyTorch (`conda`, `pip`, source): pip
 - Python version: 3.8.5
 - CUDA/cuDNN version: 11.1.105


## Additional context

Already tried it with the nightly release as well

Here you can see the full error trace;
File "E:/.../training/test.py", line 3, in <module>
    model = torch.jit.load('./models/object_detection_epochs_1_1607354042.pt')
  File "E:\...\lib\site-packages\torch\jit\_serialization.py", line 161, in load
    cpp_module = torch._C.import_ir_module(cu, f, map_location, _extra_files)
RuntimeError: 
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.
:
  File "E:\...\lib\site-packages\torchvision\ops\boxes.py", line 42
    """
    _assert_has_ops()
    return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
           ~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
Serialized   File "code/__torch__/torchvision/ops/boxes.py", line 93
  _42 = __torch__.torchvision.extension._assert_has_ops
  _43 = _42()
  _44 = ops.torchvision.nms(boxes, scores, iou_threshold)
        ~~~~~~~~~~~~~~~~~~~ <--- HERE
  return _44
'nms' is being compiled since it was called from 'batched_nms'
  File "E:\...\lib\site-packages\torchvision\ops\boxes.py", line 88
        offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes))
        boxes_for_nms = boxes + offsets[:, None]
        keep = nms(boxes_for_nms, scores, iou_threshold)
               ~~~ <--- HERE
        return keep
Serialized   File "code/__torch__/torchvision/ops/boxes.py", line 50
    _18 = torch.slice(offsets, 0, 0, 9223372036854775807, 1)
    boxes_for_nms = torch.add(boxes, torch.unsqueeze(_18, 1), alpha=1)
    keep = __torch__.torchvision.ops.boxes.nms(boxes_for_nms, scores, iou_threshold, )
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    _11 = keep
  return _11
'batched_nms' is being compiled since it was called from 'RegionProposalNetwork.filter_proposals'
Serialized   File "code/__torch__/torchvision/models/detection/rpn.py", line 64
    _11 = __torch__.torchvision.ops.boxes.clip_boxes_to_image
    _12 = __torch__.torchvision.ops.boxes.remove_small_boxes
    _13 = __torch__.torchvision.ops.boxes.batched_nms
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    num_images = (torch.size(proposals))[0]
    device = ops.prim.device(proposals)
'RegionProposalNetwork.filter_proposals' is being compiled since it was called from 'RegionProposalNetwork.forward'
  File "E:\...\lib\site-packages\torchvision\models\detection\rpn.py", line 344
        proposals = self.box_coder.decode(pred_bbox_deltas.detach(), anchors)
        proposals = proposals.view(num_images, -1, 4)
        boxes, scores = self.filter_proposals(proposals, objectness, images.image_sizes, num_anchors_per_level)
                        ~~~~~~~~~~~~~~~~~~~~~ <--- HERE

        losses = {}
Serialized   File "code/__torch__/torchvision/models/detection/rpn.py", line 37
    proposals = (self.box_coder).decode(torch.detach(pred_bbox_deltas0), anchors, )
    proposals0 = torch.view(proposals, [num_images, -1, 4])
    _8 = (self).filter_proposals(proposals0, objectness0, images.image_sizes, num_anchors_per_level, )
                                                                              ~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    boxes, scores, = _8
    losses = annotate(Dict[str, Tensor], {})


Process finished with exit code 1



cc @gmagogsfm