diff --git a/m3gnet/graph/_converters.py b/m3gnet/graph/_converters.py
index 4ae39d9..99fc373 100644
--- a/m3gnet/graph/_converters.py
+++ b/m3gnet/graph/_converters.py
@@ -8,7 +8,7 @@
 import numpy as np
 import tensorflow as tf
 from ase import Atoms
-from pymatgen.core import Structure
+from pymatgen.core.structure import Structure
 
 from m3gnet.config import DataType
 from m3gnet.type import StructureOrMolecule
diff --git a/m3gnet/graph/_structure.py b/m3gnet/graph/_structure.py
index c0a6a4c..8933aed 100644
--- a/m3gnet/graph/_structure.py
+++ b/m3gnet/graph/_structure.py
@@ -6,7 +6,7 @@
 
 import numpy as np
 from ase import Atoms
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 from pymatgen.optimization.neighbors import find_points_in_spheres
 
 from m3gnet.config import DataType
diff --git a/m3gnet/graph/tests/test_batch.py b/m3gnet/graph/tests/test_batch.py
index 69128a6..4e32811 100644
--- a/m3gnet/graph/tests/test_batch.py
+++ b/m3gnet/graph/tests/test_batch.py
@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.graph import (
     MaterialGraph,
diff --git a/m3gnet/graph/tests/test_compute.py b/m3gnet/graph/tests/test_compute.py
index c44b3e9..0ed0645 100644
--- a/m3gnet/graph/tests/test_compute.py
+++ b/m3gnet/graph/tests/test_compute.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 import tensorflow as tf
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.graph import (
     Index,
diff --git a/m3gnet/graph/tests/test_converters.py b/m3gnet/graph/tests/test_converters.py
index 0dedd83..8354d8a 100644
--- a/m3gnet/graph/tests/test_converters.py
+++ b/m3gnet/graph/tests/test_converters.py
@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.graph import BaseGraphConverter, MaterialGraph, RadiusCutoffGraphConverter
 
diff --git a/m3gnet/graph/tests/test_types.py b/m3gnet/graph/tests/test_types.py
index c9a2df4..a31b629 100644
--- a/m3gnet/graph/tests/test_types.py
+++ b/m3gnet/graph/tests/test_types.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 import tensorflow as tf
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.graph import Index, MaterialGraph, RadiusCutoffGraphConverter
 
diff --git a/m3gnet/layers/_atom_ref.py b/m3gnet/layers/_atom_ref.py
index 87f26bd..0b265d3 100644
--- a/m3gnet/layers/_atom_ref.py
+++ b/m3gnet/layers/_atom_ref.py
@@ -8,7 +8,8 @@
 import numpy as np
 import tensorflow as tf
 from ase import Atoms
-from pymatgen.core import Element, Molecule, Structure
+from pymatgen.core.periodic_table import Element
+from pymatgen.core.structure import Molecule, Structure
 
 from m3gnet.config import DataType
 from m3gnet.graph import Index
diff --git a/m3gnet/layers/tests/test_2_3_body.py b/m3gnet/layers/tests/test_2_3_body.py
index 4cb723e..ac51ce5 100644
--- a/m3gnet/layers/tests/test_2_3_body.py
+++ b/m3gnet/layers/tests/test_2_3_body.py
@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.core.structure import Lattice, Molecule, Structure
 
 from m3gnet.graph import RadiusCutoffGraphConverter, tf_compute_distance_angle
 from m3gnet.layers import PairDistance, PairVector, SphericalBesselWithHarmonics
diff --git a/m3gnet/layers/tests/test_atom_ref.py b/m3gnet/layers/tests/test_atom_ref.py
index 6ecfa9a..e51fc3a 100644
--- a/m3gnet/layers/tests/test_atom_ref.py
+++ b/m3gnet/layers/tests/test_atom_ref.py
@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.layers import AtomRef
 
diff --git a/m3gnet/layers/tests/test_bond.py b/m3gnet/layers/tests/test_bond.py
index f599a6f..31811b0 100644
--- a/m3gnet/layers/tests/test_bond.py
+++ b/m3gnet/layers/tests/test_bond.py
@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.core.structure import Lattice, Molecule, Structure
 
 from m3gnet.graph import Index, RadiusCutoffGraphConverter
 from m3gnet.layers import PairRadialBasisExpansion
diff --git a/m3gnet/layers/tests/test_core.py b/m3gnet/layers/tests/test_core.py
index 919bfe4..b8661cf 100644
--- a/m3gnet/layers/tests/test_core.py
+++ b/m3gnet/layers/tests/test_core.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 import tensorflow as tf
-from pymatgen.core import Lattice, Structure
+from pymatgen.core.structure import Lattice, Structure
 
 from m3gnet.graph import RadiusCutoffGraphConverter
 from m3gnet.layers import MLP, Embedding, GatedMLP, Pipe
diff --git a/m3gnet/layers/tests/test_gn.py b/m3gnet/layers/tests/test_gn.py
index b8c9fd6..72f3d03 100644
--- a/m3gnet/layers/tests/test_gn.py
+++ b/m3gnet/layers/tests/test_gn.py
@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.core.structure import Lattice, Molecule, Structure
 
 from m3gnet.graph import MaterialGraph, RadiusCutoffGraphConverter
 from m3gnet.layers import MLP
diff --git a/m3gnet/models/_base.py b/m3gnet/models/_base.py
index 69e77f9..9e506a3 100644
--- a/m3gnet/models/_base.py
+++ b/m3gnet/models/_base.py
@@ -8,7 +8,7 @@
 from typing import List, Union
 
 import tensorflow as tf
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 
 from m3gnet.config import DataType
 from m3gnet.graph import Index, MaterialGraph, assemble_material_graph
diff --git a/m3gnet/models/_dynamics.py b/m3gnet/models/_dynamics.py
index 5a57f5f..419702f 100644
--- a/m3gnet/models/_dynamics.py
+++ b/m3gnet/models/_dynamics.py
@@ -22,7 +22,7 @@
 from ase.optimize.mdmin import MDMin
 from ase.optimize.optimize import Optimizer
 from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 from pymatgen.io.ase import AseAtomsAdaptor
 
 from ._base import Potential
diff --git a/m3gnet/models/tests/test_model.py b/m3gnet/models/tests/test_model.py
index 3383fde..c4c2ffe 100644
--- a/m3gnet/models/tests/test_model.py
+++ b/m3gnet/models/tests/test_model.py
@@ -3,7 +3,7 @@
 
 import numpy as np
 from monty.tempfile import ScratchDir
-from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.core.structure import Lattice, Molecule, Structure
 
 from m3gnet.models import M3GNet, MolecularDynamics, Potential, Relaxer
 
diff --git a/m3gnet/trainers/_potential.py b/m3gnet/trainers/_potential.py
index 7e481ad..fa60817 100644
--- a/m3gnet/trainers/_potential.py
+++ b/m3gnet/trainers/_potential.py
@@ -6,7 +6,7 @@
 
 import tensorflow as tf
 from ase import Atoms
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 
 from m3gnet.callbacks import ManualStop
 from m3gnet.graph import Index, MaterialGraph, MaterialGraphBatchEnergyForceStress
diff --git a/m3gnet/trainers/_property.py b/m3gnet/trainers/_property.py
index 7581999..fc536e6 100644
--- a/m3gnet/trainers/_property.py
+++ b/m3gnet/trainers/_property.py
@@ -10,7 +10,7 @@
 import numpy as np
 import tensorflow as tf
 from ase import Atoms
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 
 from m3gnet.callbacks import ManualStop
 from m3gnet.graph import MaterialGraph, MaterialGraphBatch
diff --git a/m3gnet/trainers/tests/test_trainer.py b/m3gnet/trainers/tests/test_trainer.py
index 3bfe4bf..a1134e9 100644
--- a/m3gnet/trainers/tests/test_trainer.py
+++ b/m3gnet/trainers/tests/test_trainer.py
@@ -5,7 +5,7 @@
 import numpy as np
 import tensorflow as tf
 from monty.tempfile import ScratchDir
-from pymatgen.core import Structure
+from pymatgen.core.structure import Structure
 
 from m3gnet.models import M3GNet, Potential
 from m3gnet.trainers import PotentialTrainer, Trainer
diff --git a/m3gnet/type.py b/m3gnet/type.py
index 660cc37..8f073a3 100644
--- a/m3gnet/type.py
+++ b/m3gnet/type.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 from ase.atoms import Atoms
-from pymatgen.core import Molecule, Structure
+from pymatgen.core.structure import Molecule, Structure
 
 OptStrOrCallable = Optional[Union[str, Callable[..., Any]]]
 StructureOrMolecule = Union[Structure, Molecule, Atoms]
diff --git a/setup.py b/setup.py
index fde2d6e..d41e7e4 100644
--- a/setup.py
+++ b/setup.py
@@ -33,7 +33,7 @@
     name="m3gnet",
     version=version,
     install_requires=[
-        "pymatgen>=2022.1.1",
+        "pymatgen",
         "tensorflow",
         "numpy",
         "monty",