Skip to content

GitLab

Unverified Commit 4676031c authored by Cao Yuhang's avatar Cao Yuhang Committed by GitHub
Browse files

Move onnx op import inside symbolic (#453) 

* add pat

* move onnx op import to symbolic
parent 829d1e38
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
mmcv/ops/nms.py
View file @ 4676031c
...@@ -2,7 +2,6 @@ import sys ...@@ -2,7 +2,6 @@ import sys
import numpy as np import numpy as np
import torch import torch
from torch.onnx.symbolic_opset9 import select, squeeze, unsqueeze
from mmcv.utils import deprecated_api_warning from mmcv.utils import deprecated_api_warning
from ..utils import ext_loader from ..utils import ext_loader
...@@ -21,6 +20,8 @@ class NMSop(torch.autograd.Function): ...@@ -21,6 +20,8 @@ class NMSop(torch.autograd.Function):
@staticmethod @staticmethod
def symbolic(g, bboxes, scores, iou_threshold, offset): def symbolic(g, bboxes, scores, iou_threshold, offset):
from torch.onnx.symbolic_opset9 import select, squeeze, unsqueeze
boxes = unsqueeze(g, bboxes, 0) boxes = unsqueeze(g, bboxes, 0)
scores = unsqueeze(g, unsqueeze(g, scores, 0), 0) scores = unsqueeze(g, unsqueeze(g, scores, 0), 0)
max_output_per_class = g.op( max_output_per_class = g.op(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment