!6882 fix bug of arithmetic_fp16

Merge pull request !6882 from 陶云浩/bug_fix
5 years ago · f5128faba5
parent a4366b0aab 02af6670b6
commit f5128faba5
1 changed files with 2 additions and 1 deletions
--- a/mindspore/lite/src/runtime/kernel/arm/fp16/arithmetic_fp16.cc
+++ b/mindspore/lite/src/runtime/kernel/arm/fp16/arithmetic_fp16.cc
@ -156,7 +156,8 @@ int ArithmeticFP16CPUKernel::BroadcastRun(float16_t *input0, float16_t *input1,
 int ArithmeticFP16CPUKernel::DoArithmetic(int task_id) {
  int stride_per_thread = UP_DIV(param_->broadcasting_ ? outside_ : param_->out_elements_num_, context_->thread_num_);
  int cur_offset = stride_per_thread * task_id;
-  int cur_count = MSMIN(stride_per_thread, param_->out_elements_num_ - cur_offset);
+  int cur_count = param_->broadcasting_ ? MSMIN(stride_per_thread, outside_ - cur_offset)
+                                        : MSMIN(stride_per_thread, param_->out_elements_num_ - cur_offset);

  int ret = RET_OK;
  if (param_->broadcasting_) {