pt: make gelu jit happy

deepmodeling · Mar 28, 2024 · 7818956 · 7818956
1 parent 6da8eef
commit 7818956
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/deepmd/pt/utils/utils.py b/deepmd/pt/utils/utils.py
@@ -23,7 +23,12 @@ def get_activation_fn(activation: str) -> Callable:
     if activation.lower() == "relu":
         return F.relu
     elif activation.lower() == "gelu" or activation.lower() == "gelu_tf":
-        return lambda x: F.gelu(x, approximate="tanh")
+
+        @torch.jit.script
+        def gelu_tanh(x):
+            return F.gelu(x, approximate="tanh")
+
+        return gelu_tanh
     elif activation.lower() == "tanh":
         return torch.tanh
     elif activation.lower() == "relu6":