disabling torch.jit.script here for massive performance boost when us…

…ing torch.compile, our default. see issue karpathy#11. thanks @vgoklani for flagging
isbee · Jan 2, 2023 · 177d5f7 · 177d5f7
1 parent ea4de19
commit 177d5f7
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/model.py b/model.py
@@ -14,8 +14,8 @@
 import torch.nn as nn
 from torch.nn import functional as F
 
-@torch.jit.script
-def fused_gelu(x):
+# @torch.jit.script # good to enable when not using torch.compile, disable when using (our default)
+def new_gelu(x):
     """
     Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT).
     Reference: Gaussian Error Linear Units (GELU) paper: https://arxiv.org/abs/1606.08415
@@ -71,7 +71,7 @@ def __init__(self, config):
 
     def forward(self, x):
         x = self.c_fc(x)
-        x = fused_gelu(x)
+        x = new_gelu(x)
         x = self.c_proj(x)
         x = self.dropout(x)
         return x