From a43fe26a016ad7cc00183a4b7ecc50b64cf7592c Mon Sep 17 00:00:00 2001
From: zhaozhenlong <zhaozhenlong1@huawei.com>
Date: Fri, 23 Oct 2020 15:52:42 +0800
Subject: [PATCH] nms op registry

---
 mindspore/lite/src/ops/non_max_suppression.cc         | 11 +++++++++++
 .../src/ops/populate/non_max_suppression_populate.cc  |  3 ++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/mindspore/lite/src/ops/non_max_suppression.cc b/mindspore/lite/src/ops/non_max_suppression.cc
index 3faaf1b8ad..2f933cd859 100644
--- a/mindspore/lite/src/ops/non_max_suppression.cc
+++ b/mindspore/lite/src/ops/non_max_suppression.cc
@@ -16,6 +16,10 @@
 
 #include "src/ops/non_max_suppression.h"
 
+#ifndef PRIMITIVE_WRITEABLE
+#include "src/ops/ops_register.h"
+#endif
+
 namespace mindspore {
 namespace lite {
 #ifdef PRIMITIVE_WRITEABLE
@@ -44,6 +48,13 @@ int NonMaxSuppression::UnPackToFlatBuilder(const schema::Primitive *primitive, f
 int NonMaxSuppression::GetCenterPointBox() const {
   return this->primitive_->value_as_NonMaxSuppression()->centerPointBox();
 }
+
+PrimitiveC *NonMaxSuppressionCreator(const schema::Primitive *primitive) {
+  return PrimitiveC::NewPrimitiveC<NonMaxSuppression>(primitive);
+}
+
+Registry NonMaxSuppressionRegistry(schema::PrimitiveType_NonMaxSuppression, NonMaxSuppressionCreator);
+
 #endif
 int NonMaxSuppression::InferShape(std::vector<Tensor *> inputs_, std::vector<Tensor *> outputs_) {
   auto input = inputs_.front();
diff --git a/mindspore/lite/src/ops/populate/non_max_suppression_populate.cc b/mindspore/lite/src/ops/populate/non_max_suppression_populate.cc
index 4030783088..3fc60b09ca 100644
--- a/mindspore/lite/src/ops/populate/non_max_suppression_populate.cc
+++ b/mindspore/lite/src/ops/populate/non_max_suppression_populate.cc
@@ -35,7 +35,8 @@ OpParameter *PopulateNonMaxSuppressionParameter(const mindspore::lite::Primitive
   param->center_point_box_ = prim->GetCenterPointBox();
   return reinterpret_cast<OpParameter *>(param);
 }
-Registry NonMaxSuppressionParameterRegistry(schema::PrimitiveType_OneHot, PopulateNonMaxSuppressionParameter);
+Registry NonMaxSuppressionParameterRegistry(schema::PrimitiveType_NonMaxSuppression,
+                                            PopulateNonMaxSuppressionParameter);
 
 }  // namespace lite
 }  // namespace mindspore