!3363 fix cloner when funcgraph return is null

Merge pull request !3363 from leopz/fix_clone
This commit is contained in:
mindspore-ci-bot 2020-07-23 17:27:07 +08:00 committed by Gitee
commit 251fba00f5
1 changed files with 8 additions and 4 deletions

View File

@ -181,11 +181,15 @@ void Cloner::CloneFuncGraphValueNodes(const FuncGraphPtr &func_graph, const Func
MS_EXCEPTION_IF_NULL(func_graph);
MS_EXCEPTION_IF_NULL(target_func_graph);
MS_EXCEPTION_IF_NULL(manager_);
auto return_node = repl_node_[func_graph->get_return()]->cast<CNodePtr>();
if (return_node == nullptr) {
MS_LOG(EXCEPTION) << "Can't find replicate node for return.";
auto old_return = func_graph->get_return();
if (old_return != nullptr) {
auto return_node = repl_node_[old_return]->cast<CNodePtr>();
if (return_node == nullptr) {
MS_LOG(EXCEPTION) << "Can't find replicate node for return.";
}
target_func_graph->set_return(return_node);
}
target_func_graph->set_return(return_node);
auto &cnodes = func_graph->func_graph_cnodes_index();
for (auto &cnode : cnodes) {