Applied proposed change in...

Applied proposed change in !189 (comment 2814405).

src/recipes/RemoveFlatten.cpp

@@ -22,28 +22,15 @@
 
 
 namespace Aidge {
-
-    void removeFlatten(const std::set<NodePtr>& solution){
-        std::set<NodePtr> flattenNodes {};
-        for (const auto& node : solution) {
-            if (node->type() == "Flatten"){
-                printf("Flatten found.\n");
-                flattenNodes.insert(node);
-            }
-            else if (! (node->type() == "MatMul" || node->type() == "FC")){
-                AIDGE_THROW_OR_ABORT(std::runtime_error, "Node of type {} is not MatMul nor FC, an error during GraphMatching occured !", node->type());
-            }
-        }
-        GraphView::replace(flattenNodes, {});
-    }
-
     void removeFlatten(std::shared_ptr<GraphView> graphView){
         const auto matches = SinglePassGraphMatching(graphView).match(
             "(FC|MatMul)<-(Flatten)+"
         );
 
         for (const auto& solution : matches) {
-            removeFlatten(solution.graph->getNodes());
+            auto flattenNodes(solution.graph->getNodes());
+            flattenNodes.erase(solution.graph->rootNode());
+            GraphView::replace(flattenNodes, {});
         }
     }
 }