Update mlm.py

mlm.py

@@ -577,9 +577,9 @@ class BertEnergyModelForSequenceClassification(BertPreTrainedModel):
                 loss_fct = BCEWithLogitsLoss()
                 loss = loss_fct(logits, labels)
 
-        if not return_dict:
-            output = (logits, outputs.hidden_states, outputs.attentions)
-            return ((loss,) + output) if loss is not None else output
+        #if not return_dict:
+        #    output = (logits, outputs.hidden_states, outputs.attentions)
+        #    return ((loss,) + output) if loss is not None else output
 
         return SequenceClassifierOutput(
             loss=loss,