Update mlm.py

mlm.py

@@ -497,7 +497,7 @@ class BertEnergyModelForSequenceClassification(BertPreTrainedModel):
     def __init__(self, config, add_pooling_layer=True, pad_idx=None, classifier_dropout=0.1, num_labels=2):
         super().__init__(config)
         self.config = config
-
+        self.num_labels = num_labels
         self.model = BertEnergyModel(config, pad_idx=pad_idx)
         #self.lm_head = EnergyLMHead(config)