Fix import for rdkit Chem module

ecrl · Feb 13, 2025 · b0566de · b0566de
1 parent c1d0091
commit b0566de
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/graphchem/preprocessing/features.py b/graphchem/preprocessing/features.py
@@ -3,6 +3,7 @@
 
 import numpy as np
 import rdkit
+from rdkit import Chem
 import torch
 
 
@@ -221,7 +222,7 @@ def __init__(self, smiles: List[str]):
         ValueError
             If any provided SMILES string cannot be parsed by RDKit.
         """
-        mols = [rdkit.Chem.MolFromSmiles(smi) for smi in smiles]
+        mols = [Chem.MolFromSmiles(smi) for smi in smiles]
         for idx, mol in enumerate(mols):
             if mol is None:
                 raise ValueError(f"Unable to parse SMILES: {smiles[idx]}")