support compressing gelu_tf (#2957)

`gelu_tf` is equivalent to `gelu`. --------- Signed-off-by: Jinzhe Zeng <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
deepmodeling · Oct 30, 2023 · 9a4a4dd · 9a4a4dd
1 parent 6c99220
commit 9a4a4dd
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py
@@ -85,7 +85,10 @@ def __init__(
         # functype
         if activation_fn == ACTIVATION_FN_DICT["tanh"]:
             self.functype = 1
-        elif activation_fn == ACTIVATION_FN_DICT["gelu"]:
+        elif activation_fn in (
+            ACTIVATION_FN_DICT["gelu"],
+            ACTIVATION_FN_DICT["gelu_tf"],
+        ):
             self.functype = 2
         elif activation_fn == ACTIVATION_FN_DICT["relu"]:
             self.functype = 3