Update mlm.py

mlm.py

@@ -505,8 +505,9 @@ class BertEnergyModelForSequenceClassification(BertPreTrainedModel):
 
         #self.dropout = nn.Dropout(dropout)
         output_dim = config.hidden_size
-        self.norm = nn.LayerNorm(output_dim, eps=config.layer_norm_eps)
-        self.classifier = nn.Linear(output_dim, num_labels)
+        embed_dim = config.embedding_dim
+        self.norm = nn.LayerNorm(config.embedding_dim, eps=config.layer_norm_eps)
+        self.classifier = nn.Linear(config.embedding_dim, num_labels)
 
         self.post_init()