Update mlm.py

mlm.py

@@ -501,7 +501,7 @@ class BertEnergyModelForSequenceClassification(BertPreTrainedModel):
         self.model = BertEnergyModel(config, pad_idx=pad_idx)
         #self.lm_head = EnergyLMHead(config)
 
-        dropout = config.hidden_dropout_prob
+        self.dropout = config.hidden_dropout_prob
 
         #self.dropout = nn.Dropout(dropout)
         output_dim = config.hidden_size