fix_nms_ops_transformation (#11794)

* fix_nms_ops_transformation

* replace node when 5-9
This commit is contained in:
Chenhu Wang 2022-06-10 10:18:28 +08:00 committed by GitHub
parent 604dc4589c
commit 1066d4551f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -186,6 +186,10 @@ bool nms_to_nms9_callback_func(pattern::Matcher& m, pass::MatcherPass* impl) {
nms_9->set_friendly_name(root->get_friendly_name()); nms_9->set_friendly_name(root->get_friendly_name());
ngraph::copy_runtime_info(root, nms_9); ngraph::copy_runtime_info(root, nms_9);
// nms0-4 have one output, nms5/9 have 3 outputs.
if (std::dynamic_pointer_cast<ngraph::opset5::NonMaxSuppression>(root))
ngraph::replace_node(root, nms_9);
else
root->output(0).replace(nms_9->output(0)); root->output(0).replace(nms_9->output(0));
return true; return true;
} }