diff --git a/graphchem/preprocessing/features.py b/graphchem/preprocessing/features.py
index 08c9f59..fd86056 100644
--- 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]}")