Using output of hidden_state output of FlexBertModel for FlexBertForMaskedLM

modeling_flexbert.py

@@ -1140,7 +1140,7 @@ class FlexBertForMaskedLM(FlexBertPreTrainedModel):
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
         )
-
+        output=output[0]
         if self.masked_prediction and labels is not None:
             # flatten labels and output first
             labels = labels.view(-1)