!5985 fix parallel optimizer problem, assign is removed unexpectly

Merge pull request !5985 from gziyan/fix_parallel_optimizer
5 years ago · c7ff3bd7ed
parent 0c516dc749 644b680df3
commit c7ff3bd7ed
1 changed files with 6 additions and 1 deletions
--- a/mindspore/ccsrc/frontend/operator/ops_front_infer_function.cc
+++ b/mindspore/ccsrc/frontend/operator/ops_front_infer_function.cc
@ -636,7 +636,12 @@ AbstractBasePtr InferImplAssign(const AnalysisEnginePtr &, const PrimitivePtr &p
  CheckArgsSize(primitive->name(), args_spec_list, 2);

  MS_LOG(DEBUG) << "InferImplAssign " << args_spec_list[0];
-  return args_spec_list[0];
+  auto type = args_spec_list[0]->BuildType();
+  if (type->type_id() == kObjectTypeRefKey) {
+    return args_spec_list[1]->Broaden();
+  } else {
+    return args_spec_list[0];
+  }
 }

 REGISTER_FRONTENT_PRIMITIVE_EVAL_IMPL(TypeOf, prim::kPrimTypeOf, InferImplTypeof);