diff --git a/dpdata/system.py b/dpdata/system.py
index a84f31c59..bc6312f38 100644
--- a/dpdata/system.py
+++ b/dpdata/system.py
@@ -964,7 +964,12 @@ def predict(self, dp):
         labeled_sys LabeledSystem
             The labeled system.
         """
-        import deepmd.DeepPot as DeepPot
+        try:
+            # DP 1.x
+            import deepmd.DeepPot as DeepPot
+        except ModuleNotFoundError:
+            # DP 2.x
+            from deepmd.infer import DeepPot
         if not isinstance(dp, DeepPot):
             dp = DeepPot(dp)
         type_map = dp.get_type_map()
@@ -1800,7 +1805,12 @@ def from_deepmd_npy(self, folder):
         return self
 
     def predict(self, dp):
-        import deepmd.DeepPot as DeepPot
+        try:
+            # DP 1.x
+            import deepmd.DeepPot as DeepPot
+        except ModuleNotFoundError:
+            # DP 2.x
+            from deepmd.infer import DeepPot
         if not isinstance(dp, DeepPot):
             dp = DeepPot(dp)
         new_multisystems = dpdata.MultiSystems()