Fix incomplete outputs of FlaxBert (#18772)

* Fix incomplete FlaxBert outputs * fix big_bird electra roberta
huggingface · Aug 26, 2022 · 21f6f58 · 21f6f58
1 parent 62ceb4d
commit 21f6f58
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/src/transformers/models/bert/modeling_flax_bert.py b/src/transformers/models/bert/modeling_flax_bert.py
@@ -612,7 +612,7 @@ def __call__(
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
-        outputs = (hidden_states,)
+        outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions)
 
         if not return_dict:
             return tuple(v for v in outputs if v is not None)

diff --git a/src/transformers/models/big_bird/modeling_flax_big_bird.py b/src/transformers/models/big_bird/modeling_flax_big_bird.py
@@ -1438,7 +1438,7 @@ def __call__(
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
-        outputs = (hidden_states,)
+        outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions)
 
         if not return_dict:
             return tuple(v for v in outputs if v is not None)

diff --git a/src/transformers/models/electra/modeling_flax_electra.py b/src/transformers/models/electra/modeling_flax_electra.py
@@ -590,7 +590,7 @@ def __call__(
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
-        outputs = (hidden_states,)
+        outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions)
 
         if not return_dict:
             return tuple(v for v in outputs if v is not None)

diff --git a/src/transformers/models/roberta/modeling_flax_roberta.py b/src/transformers/models/roberta/modeling_flax_roberta.py
@@ -580,7 +580,7 @@ def __call__(
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
-        outputs = (hidden_states,)
+        outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions)
 
         if not return_dict:
             return tuple(v for v in outputs if v is not None)