Update flaubert2_model.py

flaubert2_model.py

@@ -390,6 +390,10 @@ class Flaubert2Model(RobertaModel):
 
         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:]