Fix FlaxBigBirdEmbeddings (#17842)

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

Fix FlaxBigBirdEmbeddings (#17842)
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
8bb2c387 · Yih-Dar · GitHub · b68d408f · 8bb2c387
Unverified Commit 8bb2c387 authored Jul 01, 2022 by Yih-Dar Committed by GitHub Jul 01, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/transformers/models/big_bird/modeling_flax_big_bird.py src/transformers/models/big_bird/modeling_flax_big_bird.py +1 -1

No files found.
--- a/src/transformers/models/big_bird/modeling_flax_big_bird.py
+++ b/src/transformers/models/big_bird/modeling_flax_big_bird.py
@@ -229,8 +229,8 @@ class FlaxBigBirdEmbeddings(nn.Module):
        hidden_states = inputs_embeds + token_type_embeddings + position_embeds

        # Layer Norm
-        hidden_states = self.LayerNorm(hidden_states)
        hidden_states = self.dropout(hidden_states, deterministic=deterministic)
+        hidden_states = self.LayerNorm(hidden_states)
        return hidden_states