Update mlm.py

mlm.py

@@ -545,8 +545,8 @@ class BertEnergyModelForSequenceClassification(BertPreTrainedModel):
         )
         
     def forward(self, input_ids, labels=None, return_dict=None, **kwargs):
-        if return_dict is None:
-            return_dict = self.return_dict
+        #if return_dict is None:
+        #    return_dict = self.return_dict
 
         outputs = self.model(input_ids, attention_mask=attention_mask, **kwargs)
         last_hidden_state = self.norm(outputs.last_hidden_state)