Update trainer.py

trainer.py

@@ -581,10 +581,10 @@ class Trainer(object):
                 drug_graphs, real_graphs, a_tensor, x_tensor, drugs_a_tensor, drugs_x_tensor, z, z_edge, z_node = bulk_data
                 
                 if self.submodel == "CrossLoss":
-                    GAN1_input_e = drugs_a_tensor
-                    GAN1_input_x = drugs_x_tensor
-                    GAN1_disc_e = a_tensor
-                    GAN1_disc_x = x_tensor
+                    GAN1_input_e = a_tensor
+                    GAN1_input_x = x_tensor
+                    GAN1_disc_e = drugs_a_tensor
+                    GAN1_disc_x = drugs_x_tensor
                 elif self.submodel == "Ligand":
                     GAN1_input_e = a_tensor
                     GAN1_input_x = x_tensor
@@ -737,11 +737,13 @@ class Trainer(object):
             drug_smiles = [line for line in open("data/chembl_train.smi", 'r').read().splitlines()]
         else:
             drug_smiles = [line for line in open("data/akt_train.smi", 'r').read().splitlines()]
-        
-        drug_mols = [Chem.MolFromSmiles(smi) for smi in drug_smiles]
-        drug_scaf = [MurckoScaffold.GetScaffoldForMol(x) for x in drug_mols]
-        fps_r = [Chem.RDKFingerprint(x) for x in drug_scaf]
 
+        if self.submodel == "RL":
+            drug_mols = [Chem.MolFromSmiles(smi) for smi in drug_smiles]
+            drug_scaf = [MurckoScaffold.GetScaffoldForMol(x) for x in drug_mols]
+            fps_r = [Chem.RDKFingerprint(x) for x in drug_scaf]
+        else:
+            fps_r = None
         akt1_human_adj = torch.load("data/akt/AKT1_human_adj.pt").reshape(1,-1).to(self.device).float() 
         akt1_human_annot = torch.load("data/akt/AKT1_human_annot.pt").reshape(1,-1).to(self.device).float() 
         
@@ -819,10 +821,6 @@ class Trainer(object):
                     GAN1_input_x = x_tensor
                     GAN1_disc_e = drugs_a_tensor
                     GAN1_disc_x = drugs_x_tensor
-                    GAN2_input_e = drugs_a_tensor
-                    GAN2_input_x = drugs_x_tensor
-                    GAN2_disc_e = a_tensor
-                    GAN2_disc_x = x_tensor
                 elif self.submodel == "Ligand":
                     GAN1_input_e = a_tensor
                     GAN1_input_x = x_tensor