[Flax] Fix BigBird (#13380)

* finish * finish

[Flax] Fix BigBird (#13380)
* finish * finish
4475f1dc · Patrick von Platen · GitHub · ecd53971 · 4475f1dc
Unverified Commit 4475f1dc authored Sep 01, 2021 by Patrick von Platen Committed by GitHub Sep 01, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

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

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
@@ -2029,6 +2029,7 @@ class FlaxBigBirdForQuestionAnswering(FlaxBigBirdPreTrainedModel):
            if token_type_ids is None:
                token_type_ids = (~logits_mask).astype("i4")
            logits_mask = jnp.expand_dims(logits_mask, axis=2)
+            logits_mask = jax.ops.index_update(logits_mask, jax.ops.index[:, 0], False)

        # init input tensors if not passed
        if token_type_ids is None: