Update modeling_protst.py

modeling_protst.py

@@ -55,6 +55,8 @@ class BertForPubMed(BertPreTrainedModel):
         self.text_mlp = ProtSTHead(config)
         self.word_mlp = ProtSTHead(config)
 
+        self.post_init() # NOTE
+
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,
@@ -111,7 +113,7 @@ class EsmForProteinRepresentation(EsmPreTrainedModel):
         self.protein_mlp = ProtSTHead(config)
         self.residue_mlp = ProtSTHead(config)
 
-        self.init_weights()
+        self.post_init() # NOTE
 
     def forward(
         self,
@@ -163,6 +165,8 @@ class EsmForProteinPropertyPrediction(EsmPreTrainedModel):
         self.model = EsmForProteinRepresentation(config)
         self.classifier = ProtSTHead(config, out_dim=config.num_labels)
 
+        self.post_init() # NOTE
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,