Refactor decode_smiles and stage connection

pluskal-lab · Aug 14, 2024 · 22f788b · 22f788b
1 parent 49afb71
commit 22f788b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/massspecgym/models/de_novo/smiles_tranformer.py b/massspecgym/models/de_novo/smiles_tranformer.py
@@ -107,9 +107,9 @@ def step(self, batch: dict, stage: Stage = Stage.NONE) -> dict:
 
         # Generate SMILES strings
         if stage in self.log_only_loss_at_stages:
-            mols_pred = self.decode_smiles(batch["spec"])
-        else:
             mols_pred = None
+        else:
+            mols_pred = self.decode_smiles(batch["spec"])
 
         return dict(loss=loss, mols_pred=mols_pred)