Update mlm.py

mlm.py

@@ -494,7 +494,7 @@ class BertEnergyModelForSequenceClassification(BertPreTrainedModel):
     ignore_index = -100
     _tied_weights_keys = ["lm_head.decoder.weight"]
 
-    def __init__(self, config, add_pooling_layer=True, pad_idx=None):
+    def __init__(self, config, add_pooling_layer=True, pad_idx=None, classifier_dropout=0.1, num_labels=2):
         super().__init__(config)
         self.config = config