Update train.py

train.py

@@ -161,7 +161,7 @@ class markerModel(pl.LightningModule):
         # self.sigmoid = nn.Sigmoid()
 
         #-- Pretrained Model Setting
-        acc_config = AutoConfig.from_pretrained("seyonec/SMILES_BPE_PubChem_100k_shard00")
+        acc_config = AutoConfig.from_pretrained(acc_model_name)
         if d_pretrained is False:
             self.d_model = RobertaModel(acc_config)
             print('acceptor model without pretraining')
@@ -170,7 +170,7 @@ class markerModel(pl.LightningModule):
                                                         output_hidden_states=True,
                                                         output_attentions=True)
         
-        don_config = AutoConfig.from_pretrained("seyonec/SMILES_BPE_PubChem_100k_shard00")
+        don_config = AutoConfig.from_pretrained(don_model_name)
 
         if p_pretrained is False:
             self.p_model = RobertaModel(don_config)