Update flaubert2_model.py

flaubert2_model.py

@@ -388,6 +388,9 @@ 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
 
         if not return_dict:
@@ -397,7 +400,7 @@ class Flaubert2Model(RobertaModel):
             last_hidden_state=sequence_output,
             pooler_output=pooled_output,
             past_key_values=encoder_outputs.past_key_values,
-            hidden_states=encoder_outputs.hidden_states,
+            hidden_states=hidden_states,
             attentions=encoder_outputs.attentions,
             cross_attentions=encoder_outputs.cross_attentions,
         )