diff --git a/quantumml/featurizers/base_cleasses.py b/quantumml/featurizers/base_cleasses.py
index 875eee4..e6276f8 100644
--- a/quantumml/featurizers/base_cleasses.py
+++ b/quantumml/featurizers/base_cleasses.py
@@ -178,7 +178,7 @@ def featurize(self,
       `structures`.
     """
     try:
-      from pymatgen import Structure
+      from pymatgen.core import Structure
     except ModuleNotFoundError:
       raise ImportError("This class requires pymatgen to be installed.")
 
diff --git a/quantumml/mlModels.py b/quantumml/mlModels.py
index deb4df3..757be4e 100644
--- a/quantumml/mlModels.py
+++ b/quantumml/mlModels.py
@@ -8,7 +8,7 @@
 import ast
 import base64
 from dscribe.descriptors import SOAP
-from pymatgen import Structure
+from pymatgen.core import Structure
 class MLModel:
     """
     Base class for pre-trained machine learning model