Upload model

transformer_fnqs.py

@@ -57,8 +57,8 @@ class EncoderBlock(nn.Module):
         self.layer_norm_2 = nn.LayerNorm(dtype=jnp.float64, param_dtype=jnp.float64)
 
         self.ff = nn.Sequential([
-            nn.Dense(2*self.d_model, kernel_init=nn.initializers.xavier_uniform(), param_dtype=jnp.float64, dtype=jnp.float64),
-            nn.relu,
+            nn.Dense(4*self.d_model, kernel_init=nn.initializers.xavier_uniform(), param_dtype=jnp.float64, dtype=jnp.float64),
+            nn.gelu,
             nn.Dense(self.d_model, kernel_init=nn.initializers.xavier_uniform(), param_dtype=jnp.float64, dtype=jnp.float64),
         ])