From 1957afa243b1af7a90835241ed4e819c3fbef846 Mon Sep 17 00:00:00 2001
From: liuwenhao4 <liuwenhao4@huawei.com>
Date: Wed, 9 Dec 2020 17:53:39 +0800
Subject: [PATCH] add NHWC format requirement of topk ops

---
 mindspore/lite/src/ops/topk.cc           | 4 ++++
 mindspore/lite/tools/common/node_util.cc | 3 ++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/mindspore/lite/src/ops/topk.cc b/mindspore/lite/src/ops/topk.cc
index 2c51c19cf22..24eab142c3b 100644
--- a/mindspore/lite/src/ops/topk.cc
+++ b/mindspore/lite/src/ops/topk.cc
@@ -60,6 +60,10 @@ int TopK::InferShape(std::vector<Tensor *> inputs_, std::vector<Tensor *> output
   }
   auto input = inputs_.front();
   MS_ASSERT(input != nullptr);
+  if (input->format() != schema::Format::Format_NHWC) {
+    MS_LOG(ERROR) << "topk only support NHWC now!";
+    return RET_FORMAT_ERR;
+  }
   auto output0 = outputs_.front();
   MS_ASSERT(output0 != nullptr);
   auto output1 = outputs_.at(1);
diff --git a/mindspore/lite/tools/common/node_util.cc b/mindspore/lite/tools/common/node_util.cc
index 29b4d5d0fde..efa1574e106 100644
--- a/mindspore/lite/tools/common/node_util.cc
+++ b/mindspore/lite/tools/common/node_util.cc
@@ -50,7 +50,8 @@ static const std::vector<schema::PrimitiveType> nhwcOpList = {
   schema::PrimitiveType_BiasAdd,
   schema::PrimitiveType_InstanceNorm,
   schema::PrimitiveType_SpaceToDepth,
-  schema::PrimitiveType_DepthToSpace};
+  schema::PrimitiveType_DepthToSpace,
+  schema::PrimitiveType_TopK};
 
 static const std::vector<schema::PrimitiveType> nhwcOpAllInputList = {
 #ifdef SUPPORT_TRAIN