modified: ge/graph/passes/infershape_pass.cc

4 years ago · 848236b21c
parent c293465b6c
commit 848236b21c
1 changed files with 3 additions and 3 deletions
--- a/ge/graph/passes/infershape_pass.cc
+++ b/ge/graph/passes/infershape_pass.cc
@ -42,9 +42,9 @@ Status InferShapePass::Run(NodePtr &node) {
    GELOGE(GE_GRAPH_INFERSHAPE_FAILED, "infershape failed. node: %s", node->GetName().c_str());
    return GE_GRAPH_INFERSHAPE_FAILED;
  }
-  if (node->GetType() == WHILE) {
-    bool need_repass = false;
-    AttrUtils::GetBool(node->GetOpDesc(), "need_infer_again_", need_repass);
+  bool need_repass = false;
+  auto has_attr = AttrUtils::GetBool(node->GetOpDesc(), "need_infer_again_", need_repass);
+  if (has_attr) {
    if (!OptionExists(kOptimizeAfterSubGraph)) {
      return SUCCESS;
    }