Skip to content

GitLab

Commit bb7c638f authored by Boris Fomitchev's avatar Boris Fomitchev
Browse files

Restored @script for openai_gelu 


Signed-off-by: default avatarBoris Fomitchev <bfomitchev@nvidia.com>
parent fe794c5a
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
megatron/model/utils.py
View file @ bb7c638f
...@@ -48,6 +48,7 @@ def get_linear_layer(rows, columns, init_method): ...@@ -48,6 +48,7 @@ def get_linear_layer(rows, columns, init_method):
layer.bias.zero_() layer.bias.zero_()
return layer return layer
@torch.jit.script
def gelu_impl(x): def gelu_impl(x):
"""OpenAI's gelu implementation.""" """OpenAI's gelu implementation."""
return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x *
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment