Change remove_flatten to take in a node instead of a graphView

python_binding/recipies/pybind_Recipies.cpp

@@ -21,7 +21,7 @@ namespace py = pybind11;
 namespace Aidge {
 void init_Recipies(py::module &m) {
   m.def("fuse_mul_add", &fuseMulAdd, py::arg("nodes"));
-  m.def("remove_flatten", &removeFlatten, py::arg("view"));
+  m.def("remove_flatten", &removeFlatten, py::arg("nodes"));
   
 }
 } // namespace Aidge

src/recipies/RemoveFlatten.cpp

@@ -16,13 +16,9 @@
 #include "aidge/utils/Recipies.hpp"
 
 namespace Aidge {
-    void removeFlatten(std::shared_ptr<GraphView> view) {
-        for (auto& nodePtr : view->getNodes()) {
-            if (nodePtr->type() == "Flatten") {
-                auto g = std::make_shared<GraphView>();
-                g->add(std::set<std::shared_ptr<Node>>({nodePtr}));
-                g->replaceWith({});
-            }
-        }
+    void removeFlatten(std::set<std::shared_ptr<Node>> nodes) {
+        auto g = std::make_shared<GraphView>();
+        g->add(std::set<std::shared_ptr<Node>>({nodes}));
+        g->replaceWith({});
     }
 }
\ No newline at end of file