Try remove eager scripting calls (#2248) (#2362)

* Try remove eager scripting calls

* remove script call

Co-authored-by: eellison <[email protected]>
Co-authored-by: Francisco Massa <[email protected]>

Co-authored-by: eellison <[email protected]>
Co-authored-by: eellison <[email protected]>

Author: 3 people

torchvision/models/detection/_utils.py

@@ -75,7 +75,7 @@ def __call__(self, matched_idxs):
         return pos_idx, neg_idx
 
 
-@torch.jit.script
+@torch.jit._script_if_tracing
 def encode_boxes(reference_boxes, proposals, weights):
     # type: (torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor
     """

torchvision/models/detection/roi_heads.py

@@ -205,7 +205,7 @@ def _onnx_heatmaps_to_keypoints(maps, maps_i, roi_map_width, roi_map_height,
     return xy_preds_i, end_scores_i
 
 
-@torch.jit.script
+@torch.jit._script_if_tracing
 def _onnx_heatmaps_to_keypoints_loop(maps, rois, widths_ceil, heights_ceil,
                                      widths, heights, offset_x, offset_y, num_keypoints):
     xy_preds = torch.zeros((0, 3, int(num_keypoints)), dtype=torch.float32, device=maps.device)
@@ -451,7 +451,7 @@ def _onnx_paste_mask_in_image(mask, box, im_h, im_w):
     return im_mask
 
 
-@torch.jit.script
+@torch.jit._script_if_tracing
 def _onnx_paste_masks_in_image_loop(masks, boxes, im_h, im_w):
     res_append = torch.zeros(0, im_h, im_w)
     for i in range(masks.size(0)):

torchvision/ops/boxes.py

@@ -4,7 +4,6 @@
 import torchvision
 
 
-@torch.jit.script
 def nms(boxes, scores, iou_threshold):
     # type: (Tensor, Tensor, float) -> Tensor
     """
@@ -41,7 +40,7 @@ def nms(boxes, scores, iou_threshold):
     return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
 
 
-@torch.jit.script
+@torch.jit._script_if_tracing
 def batched_nms(boxes, scores, idxs, iou_threshold):
     # type: (Tensor, Tensor, Tensor, float) -> Tensor
     """