From 36d561fc5cc776b97e112add0cf5ea223120e475 Mon Sep 17 00:00:00 2001
From: AlexHJH <hjhalex2600@outlook.com>
Date: Wed, 2 Aug 2023 00:21:23 +0800
Subject: [PATCH] fix missing functions

Signed-off-by: AlexHJH <hjhalex2600@outlook.com>
---
 onnxoptimizer/optimize.h | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/onnxoptimizer/optimize.h b/onnxoptimizer/optimize.h
index f25e6c3f6..9404b8761 100644
--- a/onnxoptimizer/optimize.h
+++ b/onnxoptimizer/optimize.h
@@ -46,12 +46,23 @@ struct Optimizer {
     ModelProto mp_out = PrepareOutput(mp_in);
     this->pass_manager->run(*g);
     ExportModelProto(&mp_out, g);
+
+    // Maybe we can optimize these functions, now just copy
+    AddFunctionsToModel(mp_in, mp_out);
     return mp_out;
   }
 
  private:
   std::shared_ptr<PassManager> pass_manager;
 
+  void AddFunctionsToModel(const ModelProto &original_model,
+                           ModelProto &output_model) {
+    for (const auto& function_proto : original_model.functions()) {
+      auto* p_f = output_model.add_functions();
+      p_f->CopyFrom(function_proto);
+    }
+  }
+
   ModelProto AddInitializerToInput(const ModelProto &original_model) {
     ModelProto model = original_model;
     std::vector<std::string> input_names;