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 inputs_, std::vector 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 nhwcOpList = { schema::PrimitiveType_BiasAdd, schema::PrimitiveType_InstanceNorm, schema::PrimitiveType_SpaceToDepth, - schema::PrimitiveType_DepthToSpace}; + schema::PrimitiveType_DepthToSpace, + schema::PrimitiveType_TopK}; static const std::vector nhwcOpAllInputList = { #ifdef SUPPORT_TRAIN