Update flaubert2_model.py

flaubert2_model.py

@@ -388,11 +388,12 @@ class Flaubert2Model(RobertaModel):
 
         sequence_output = encoder_outputs[0].transpose(0,1)
 
-        # Fairseq Linformer implementation works with transposed hidden states -> we transpose them back for HF implementation.
-        hidden_states = [h.transpose(0,1) for h in encoder_outputs.hidden_states]
-        
         pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
 
+        # Fairseq Linformer implementation works with transposed hidden states -> we transpose them back for HF implementation.
+        if output_hidden_states:
+            encoder_outputs.hidden_states = [h.transpose(0,1) for h in encoder_outputs.hidden_states]
+
         if not return_dict:
             return (sequence_output, pooled_output) + encoder_outputs[1:]
 
@@ -400,7 +401,7 @@ class Flaubert2Model(RobertaModel):
             last_hidden_state=sequence_output,
             pooler_output=pooled_output,
             past_key_values=encoder_outputs.past_key_values,
-            hidden_states=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
             attentions=encoder_outputs.attentions,
             cross_attentions=encoder_outputs.cross_attentions,
         )