Fixed element type for constants in the conversion of NMS-5 to NMSIE3.

vgavrilo · Oct 16, 2020 · 7597f30 · 7597f30
1 parent 27b2cd0
commit 7597f30
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/...e/src/legacy_api/src/transformations/convert_opset1_to_legacy/convert_nms_5_to_legacy.cpp b/...e/src/legacy_api/src/transformations/convert_opset1_to_legacy/convert_nms_5_to_legacy.cpp
@@ -75,17 +75,17 @@ ngraph::pass::ConvertNMS5ToLegacyMatcher::ConvertNMS5ToLegacyMatcher() {
         }
         auto new_iou_threshold = arg3;
         if (iou_threshold_rank.get_length() == 0) {
-            new_iou_threshold = std::make_shared<ngraph::op::Unsqueeze>(arg3, opset1::Constant::create(element::f32, Shape{1}, {0.0f}));
+            new_iou_threshold = std::make_shared<ngraph::op::Unsqueeze>(arg3, opset1::Constant::create(element::i64, Shape{1}, {0}));
             new_ops.push_back(new_iou_threshold.get_node_shared_ptr());
         }
         auto new_score_threshold = arg4;
         if (score_threshold_rank.get_length() == 0) {
-            new_score_threshold = std::make_shared<ngraph::op::Unsqueeze>(arg4, opset1::Constant::create(element::f32, Shape{1}, {0.0f}));
+            new_score_threshold = std::make_shared<ngraph::op::Unsqueeze>(arg4, opset1::Constant::create(element::i64, Shape{1}, {0}));
             new_ops.push_back(new_score_threshold.get_node_shared_ptr());
         }
         auto new_soft_nms_sigma = arg5;
         if (soft_nms_sigma_rank.get_length() == 0) {
-            new_soft_nms_sigma = std::make_shared<ngraph::op::Unsqueeze>(arg5, opset1::Constant::create(element::f32, Shape{1}, {0.0f}));
+            new_soft_nms_sigma = std::make_shared<ngraph::op::Unsqueeze>(arg5, opset1::Constant::create(element::i64, Shape{1}, {0}));
             new_ops.push_back(new_soft_nms_sigma.get_node_shared_ptr());
         }
         int center_point_box = 0;