Fixed implementation error leading to only 1 MLP layer instead of 2 a…

…nd normalization layer
Reytuag · Nov 2, 2024 · db3e52a · db3e52a
1 parent 0f9778a
commit db3e52a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformerXL.py b/transformerXL.py
@@ -75,7 +75,7 @@ def __call__(self, values_keys:jnp.ndarray, queries:jnp.ndarray, pos_embed:jnp.n
         out = self.dense1(out_attention_n)
         out = nn.activation.gelu(out)
         #out = nn.activation.relu(out)
-        out = self.dense2(out_attention)
+        out = self.dense2(out)
         if(self.gating):
             out= self.gate2(out,jax.nn.relu(out_attention))
         else: