Unverified Commit 5df759ad authored by RunningLeon's avatar RunningLeon Committed by GitHub
Browse files

[Fix]: add header file for std::iota(#803) (#820)

* add header to std::iota

* fix lint

* fix lint for cpp
parent 459d728c
...@@ -6,6 +6,7 @@ ...@@ -6,6 +6,7 @@
#include <cmath> #include <cmath>
#include <iostream> #include <iostream>
#include <iterator> #include <iterator>
#include <numeric> // std::iota
#include <vector> #include <vector>
#include "../ort_mmcv_utils.h" #include "../ort_mmcv_utils.h"
......
...@@ -291,7 +291,8 @@ def batched_nms(boxes, scores, idxs, nms_cfg, class_agnostic=False): ...@@ -291,7 +291,8 @@ def batched_nms(boxes, scores, idxs, nms_cfg, class_agnostic=False):
nms_op = eval(nms_type) nms_op = eval(nms_type)
split_thr = nms_cfg_.pop('split_thr', 10000) split_thr = nms_cfg_.pop('split_thr', 10000)
if boxes_for_nms.shape[0] < split_thr: # Won't split to multiple nms nodes when exporting to onnx
if boxes_for_nms.shape[0] < split_thr or torch.onnx.is_in_onnx_export():
dets, keep = nms_op(boxes_for_nms, scores, **nms_cfg_) dets, keep = nms_op(boxes_for_nms, scores, **nms_cfg_)
boxes = boxes[keep] boxes = boxes[keep]
scores = dets[:, -1] scores = dets[:, -1]
......
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