use tanh approximation for gelu

algorithmic_efficiency/workloads/librispeech_conformer/librispeech_pytorch/models.py

@@ -3,6 +3,7 @@
 """
 
 from dataclasses import dataclass
+from functools import partial
 import math
 from typing import Tuple
 
@@ -213,7 +214,8 @@ def forward(self, inputs, padding_mask):
     if self.config.activation_function_name == 'swish':
       activation_fn = F.silu
     elif self.config.activation_function_name == 'gelu':
-      activation_fn = F.gelu
+      # Use tanh approximation of GELU which is default for jax
+      activation_fn = partial(F.gelu, approximate='tanh')
     else:
       raise ValueError(
           'Only "swish" and "gelu" are supported '