Fix FlaxBigBirdEmbeddings (#17842)

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
huggingface · Jul 1, 2022 · 8bb2c38 · 8bb2c38
1 parent b68d408
commit 8bb2c38
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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 @@ def __call__(self, input_ids, token_type_ids, position_ids, attention_mask, dete
         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