diff --git a/ge/graph/passes/attach_stream_label_pass.cc b/ge/graph/passes/attach_stream_label_pass.cc index 7b9ae20b..fe8d0e7d 100644 --- a/ge/graph/passes/attach_stream_label_pass.cc +++ b/ge/graph/passes/attach_stream_label_pass.cc @@ -51,13 +51,15 @@ Status AttachStreamLabelPass::ClearStatus() { /// void AttachStreamLabelPass::FindNodes(const ComputeGraphPtr &graph) { for (const NodePtr &node : graph->GetDirectNode()) { - const std::string &type = node->GetType(); - if (type == STREAMSWITCH) { + const auto &op_desc = node->GetOpDesc(); + if (op_desc == nullptr) { + continue; + } + const std::string &type = op_desc->GetType(); + if ((type == STREAMSWITCH) && op_desc->HasAttr(ATTR_NAME_SWITCH_TRUE_BRANCH_FLAG)) { stream_switch_nodes_.emplace_back(node); - } else if (type == STREAMMERGE) { - if ((node->GetOpDesc() != nullptr) && !node->GetOpDesc()->HasAttr(ATTR_NAME_NEXT_ITERATION)) { - need_label_nodes_.emplace_back(node); - } + } else if ((type == STREAMMERGE)&& !op_desc->HasAttr(ATTR_NAME_NEXT_ITERATION)) { + need_label_nodes_.emplace_back(node); } else if ((type == ENTER) || (type == REFENTER)) { enter_nodes_.emplace_back(node); } @@ -220,7 +222,7 @@ Status AttachStreamLabelPass::SetEnterLabel(const std::vector &enter_no GE_CHECK_NOTNULL(active_node); (void)AttrUtils::GetStr(active_node->GetOpDesc(), ATTR_NAME_STREAM_LABEL, stream_label); if (stream_label.empty()) { - GELOGD("stream_label of enter_active & enter_nodes is empty."); + GELOGD("stream_label of enter_active %s is empty.", active_node->GetName().c_str()); return SUCCESS; }