Imported Upstream version 1.25.0
[platform/core/ml/nnfw.git] / compiler / circle-opselector / src / OpSelector.cpp
index 31cbf9c..09a6654 100644 (file)
@@ -156,6 +156,10 @@ std::unique_ptr<loco::Graph> make_graph(const std::vector<const luci::CircleNode
       }
     }
   }
+
+  const auto original_graph = nodes.at(0)->graph();
+  const auto original_outputs = loco::output_nodes(const_cast<loco::Graph *>(original_graph));
+
   // set graph output
   for (auto &n : nodes)
   {
@@ -169,8 +173,19 @@ std::unique_ptr<loco::Graph> make_graph(const std::vector<const luci::CircleNode
         break;
       }
     }
+
+    bool originalOutput = false;
+    for (const auto &o : outputs)
+    {
+      if (std::find(original_outputs.begin(), original_outputs.end(), o) != original_outputs.end())
+      {
+        originalOutput = true;
+        break;
+      }
+    }
+
     // the node isn't graph output if it is an other node's output
-    if (beingUsed)
+    if (beingUsed and not originalOutput)
       continue;
 
     IsMultiOutputNode multiout_visitor;
@@ -335,7 +350,7 @@ OpSelector::select_by<SelectType::NAME>(const std::vector<std::string> &tokens)
     auto cnode = loco::must_cast<const luci::CircleNode *>(node);
     std::string node_name = cnode->name();
 
-    for (auto selected_name : tokens)
+    for (const auto &selected_name : tokens)
       if (selected_name.compare(node_name) == 0) // find the selected name
         selected_nodes.emplace_back(cnode);
   }