From efb831fc893cf2c76927d073ef15c9786b1ce931 Mon Sep 17 00:00:00 2001
From: zhongligeng <zhongligeng@huawei.com>
Date: Tue, 18 Aug 2020 09:35:26 +0800
Subject: [PATCH] fix space to depth

---
 mindspore/lite/src/ops/space_to_depth.cc | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mindspore/lite/src/ops/space_to_depth.cc b/mindspore/lite/src/ops/space_to_depth.cc
index e0df66b25f..79f3f175b3 100644
--- a/mindspore/lite/src/ops/space_to_depth.cc
+++ b/mindspore/lite/src/ops/space_to_depth.cc
@@ -58,9 +58,9 @@ int SpaceToDepth::InferShape(std::vector<lite::tensor::Tensor *> inputs, std::ve
   }
 
   int32_t block_size = GetBlockSize();
-  if (input_shape[NHWC_C] % (block_size * block_size) != 0 || input_shape[NHWC_C] == 0) {
-    MS_LOG(ERROR) << "input dimension c size " << input_shape[NHWC_C] << " should be mulitple of block_size("
-                  << block_size << ") * block_size)!";
+  if (input_shape[NHWC_H] % block_size != 0 || input_shape[NHWC_H] == 0 || input_shape[NHWC_W] % block_size != 0 ||
+      input_shape[NHWC_W] == 0) {
+    MS_LOG(ERROR) << "input dimension h or w size error!";
     return 1;
   }
   std::vector<int32_t> output_shape(input_shape.size());