Fix FlaxBigBirdEmbeddings (huggingface#17842)

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
viclzhu · Jul 18, 2022 · d232d4a · d232d4a
1 parent 711a4f8
commit d232d4a
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