fix bug of op space_to_depth

4 years ago · cf52dbb46d
parent 3fe8916afa
commit cf52dbb46d
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;