Repository: pytorch
Issue: Error loading TorchScript model with torchvision::nms operation in libtorch
Labels: oncall: jit

### 🐛 Describe the bug

When using libtorch 1.11.0 (debug version) to load a TorchScript model (MaskRCNN) that includes the torchvision::nms operation, an error occurs during the model loading process, indicating that the torchvision::nms operation is unknown or not supported in libtorch.

**Error Message:**
```
Error loading the model:
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 "/home/fanuc/fanuc/anaconda3/envs/pytorch_env/lib/python3.10/site-packages/torchvision/ops/boxes.py", line 40
        _log_api_usage_once(nms)
    _assert_has_ops()
    return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
           ~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
Serialized   File "code/__torch__/torchvision/ops/boxes.py", line 154
  _64 = __torch__.torchvision.extension._assert_has_ops
  _65 = _64()
  _66 = ops.torchvision.nms(boxes, scores, iou_threshold)
        ~~~~~~~~~~~~~~~~~~~ <--- HERE
  return _66
'nms' is being compiled since it was called from '_batched_nms_vanilla'
  File "/home/fanuc/fanuc/anaconda3/envs/pytorch_env/lib/python3.10/site-packages/torchvision/ops/boxes.py", line 108
    for class_id in torch.unique(idxs):
        curr_indices = torch.where(idxs == class_id)[0]
        curr_keep_indices = nms(boxes[curr_indices], scores[curr_indices], iou_threshold)
                            ~~~ <--- HERE
        keep_mask[curr_indices[curr_keep_indices]] = True
    keep_indices = torch.where(keep_mask)[0]
Serialized   File "code/__torch__/torchvision/ops/boxes.py", line 83
    _31 = torch.index(boxes, _30)
    _32 = annotate(List[Optional[Tensor]], [curr_indices])
    curr_keep_indices = __torch__.torchvision.ops.boxes.nms(_31, torch.index(scores, _32), iou_threshold, )
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    _33 = annotate(List[Optional[Tensor]], [curr_keep_indices])
    _34 = torch.index(curr_indices, _33)
'_batched_nms_vanilla' is being compiled since it was called from 'batched_nms'
Serialized   File "code/__torch__/torchvision/ops/boxes.py", line 35
    idxs: Tensor,
    iou_threshold: float) -> Tensor:
  _9 = __torch__.torchvision.ops.boxes._batched_nms_vanilla
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  _10 = __torch__.torchvision.ops.boxes._batched_nms_coordinate_trick
  _11 = torch.numel(boxes)
'batched_nms' is being compiled since it was called from 'RegionProposalNetwork.filter_proposals'
Serialized   File "code/__torch__/torchvision/models/detection/rpn.py", line 72
    _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 "/home/fanuc/fanuc/anaconda3/envs/pytorch_env/lib/python3.10/site-packages/torchvision/models/detection/rpn.py", line 353
        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 43
    proposals0 = torch.view(proposals, [num_images, -1, 4])
    image_sizes = images.image_sizes
    _8 = (self).filter_proposals(proposals0, objectness0, image_sizes, num_anchors_per_level, )
                                                                       ~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    boxes, scores, = _8
    losses = annotate(Dict[str, Tensor], {})
```

**Running Code**
```
#include <torch/script.h>   // One-stop header.
#include <torch/torch.h>
#include <opencv2/opencv.hpp>

#include <iostream>
#include <memory>

std::string model_path = "../../models/script_model.pt";
std::string image_path = "../../images/1_color.png";

int main() {

  torch::jit::script::Module module;
  try {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    torch::Device device(torch::kCPU);
    module = torch::jit::load(model_path, device);
  }
  catch (const torch::jit::ErrorReport& e) {
    std::cerr << "Error loading the model:" << e.what() << std::endl;
    return -1;
  }
  return 0
}
```

**Model and convert script**
```
def get_model_instance_segmentation(num_classes, load_pretrain_weights=True):
    # get maskrcnn model form torchvision.models
    # load an instance segmentation model pre-trained on COCO
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=load_pretrain_weights)

    # get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # now get the number of input features for the mask classifier
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # and replace the mask predictor with a new one
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                       hidden_layer,
                                                       num_classes)
    return model

def generate_torch_script(pytorch_model, device, img_file, save_path):
    # input of the model (from pil image to tensor, do not normalize image)
    original_img = Image.open(img_file).convert('RGB')      # load image
    data_transform = transforms.Compose([transforms.ToTensor()])
    img = data_transform(original_img)
    img = torch.unsqueeze(img, dim=0).to(device)        # expand batch dimension to device

    # export the model
    pytorch_model.eval()
    if device.type == 'cpu':
        pytorch_model = pytorch_model.cpu()
    traced_script_module = torch.jit.script(pytorch_model, img)
    # traced_script_module = torch.jit.trace(pytorch_model, img)
    traced_script_module.save(save_path)
```


### Versions

**Environment**
* PyTorch version: Pytorch 1.11.0
* TorchVision version: 0.12.0
* LibTorch version:  libtorch cpu debug version 1.11.0
* Operating system: Windows 10, compiler MSVC

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel