diff --git a/prot2mol/utils.py b/prot2mol/utils.py index 3f5dc1e..530f1ac 100644 --- a/prot2mol/utils.py +++ b/prot2mol/utils.py @@ -179,8 +179,16 @@ def to_mol(smiles_list): return [Chem.MolFromSmiles(smiles) for smiles in smiles_list] def sascorer_calculation(mols): - return [sascorer.calculateScore(mol) if mol is not None else None for mol in mols] - + scores = [] + for mol in mols: + if mol is None: + scores.append(None) + else: + try: + scores.append(sascorer.calculateScore(mol)) + except ZeroDivisionError: + scores.append(None) + return scores def qed_calculation(mols): return [QED.qed(mol) if mol is not None else None for mol in mols] @@ -271,4 +279,4 @@ def metrics_calculation(predictions, references, train_data, train_vec=None,trai - \ No newline at end of file +