Fix missing approximate parameters of nn.GELU

senli123 · May 9, 2024 · 53299c5 · 53299c5
1 parent 8994bd0
commit 53299c5
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/tools/pnnx/src/ir.cpp b/tools/pnnx/src/ir.cpp
@@ -2647,7 +2647,7 @@ int insert_function(FILE* pyfp, std::vector<std::string>& custom_ops_names, std:
     std::ostringstream functionContent;
     bool insideFunction = false;
 
-    std::regex functionStartRegex(R "(\s*def\s+(\w+)\s*\(.*\):)");
+    std::regex functionStartRegex(R"(\s*def\s+(\w+)\s*\(.*\):)");
     std::regex pattern("^import");
     std::smatch match;
 
@@ -2715,7 +2715,7 @@ int get_custom_op_names(std::string& customop_infer_py, std::vector<std::string>
 
     std::string line;
 
-    std::regex functionStartRegex(R "(\s*def\s+(\w+)\s*\(.*\):)");
+    std::regex functionStartRegex(R"(\s*def\s+(\w+)\s*\(.*\):)");
     std::smatch match;
 
     while (std::getline(file, line))

diff --git a/tools/pnnx/src/pass_level1/nn_GELU.cpp b/tools/pnnx/src/pass_level1/nn_GELU.cpp
@@ -13,7 +13,7 @@
 // specific language governing permissions and limitations under the License.
 
 #include "pass_level1.h"
-
+#include "../utils.h"
 namespace pnnx {
 
 class GELU : public FuseModulePass
@@ -28,6 +28,13 @@ class GELU : public FuseModulePass
     {
         return "nn.GELU";
     }
+
+    void write(Operator* op, const std::shared_ptr<torch::jit::Graph>& graph) const
+    {
+        const torch::jit::Node* celu = find_node_by_kind(graph, "aten::gelu");
+
+        op->params["approximate"] = celu->namedInput("approximate");
+    }
 };
 
 REGISTER_GLOBAL_PNNX_FUSE_MODULE_PASS(GELU)