!6363 fix bug of op space_to_depth

Merge pull request !6363 from 陶云浩/space_to_depth
4 years ago · 6d03fb822d
parent 43770c9fda cf52dbb46d
commit 6d03fb822d
1 changed files with 1 additions and 1 deletions
--- a/mindspore/lite/src/runtime/kernel/arm/fp32/space_to_depth.cc
+++ b/mindspore/lite/src/runtime/kernel/arm/fp32/space_to_depth.cc
@ -50,7 +50,7 @@ int SpaceToDepthCPUKernel::ReSize() {
    return RET_FORMAT_ERR;
  }

-  num_unit_ = static_cast<int>(in_tensors_[0]->shape().at(kNHWC_H));
+  num_unit_ = static_cast<int>(out_tensors_[0]->shape().at(kNHWC_H));
  thread_h_num_ = MSMIN(op_parameter_->thread_num_, num_unit_);
  thread_h_stride_ = UP_DIV(num_unit_, thread_h_num_);
  return RET_OK;