From 6710464f6df20537f5db32174041efc683b131c1 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 28 Aug 2025 16:55:29 +0000
Subject: [PATCH 1/6] Initial plan


From 093874157f42b4ff24260c14c37732fecde050e9 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 28 Aug 2025 17:12:16 +0000
Subject: [PATCH 2/6] feat(infer): add get_model method to DeepEval backends

---
 deepmd/dpmodel/infer/deep_eval.py    | 12 +++-
 deepmd/infer/deep_eval.py            | 28 +++++++++
 deepmd/pd/infer/deep_eval.py         | 10 +++
 deepmd/pt/infer/deep_eval.py         | 10 +++
 deepmd/tf/infer/deep_eval.py         | 10 +++
 source/tests/infer/test_get_model.py | 91 ++++++++++++++++++++++++++++
 6 files changed, 160 insertions(+), 1 deletion(-)
 create mode 100644 source/tests/infer/test_get_model.py

diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py
index 9fd96ed491..e19da62103 100644
--- a/deepmd/dpmodel/infer/deep_eval.py
+++ b/deepmd/dpmodel/infer/deep_eval.py
@@ -391,4 +391,14 @@ def _get_output_shape(self, odef, nframes, natoms):
 
     def get_model_def_script(self) -> dict:
         """Get model definition script."""
-        return json.loads(self.model.get_model_def_script())
+        return json.loads(self.dp.get_model_def_script())
+
+    def get_model(self):
+        """Get the dpmodel BaseModel.
+
+        Returns
+        -------
+        BaseModel
+            The dpmodel BaseModel.
+        """
+        return self.dp
diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py
index d067c4322e..e97f1cb0a6 100644
--- a/deepmd/infer/deep_eval.py
+++ b/deepmd/infer/deep_eval.py
@@ -341,6 +341,20 @@ def get_observed_types(self) -> dict:
         """Get observed types (elements) of the model during data statistics."""
         raise NotImplementedError("Not implemented in this backend.")
 
+    @abstractmethod
+    def get_model(self):
+        """Get the model module implemented by the deep learning framework.
+
+        For PyTorch, this returns the nn.Module. For Paddle, this returns
+        the paddle.nn.Layer. For TensorFlow, this returns the graph.
+        For dpmodel, this returns the BaseModel.
+
+        Returns
+        -------
+        model
+            The model module implemented by the deep learning framework.
+        """
+
 
 class DeepEval(ABC):
     """High-level Deep Evaluator interface.
@@ -685,3 +699,17 @@ def get_model_size(self) -> dict:
     def get_observed_types(self) -> dict:
         """Get observed types (elements) of the model during data statistics."""
         return self.deep_eval.get_observed_types()
+
+    def get_model(self):
+        """Get the model module implemented by the deep learning framework.
+
+        For PyTorch, this returns the nn.Module. For Paddle, this returns
+        the paddle.nn.Layer. For TensorFlow, this returns the graph.
+        For dpmodel, this returns the BaseModel.
+
+        Returns
+        -------
+        model
+            The model module implemented by the deep learning framework.
+        """
+        return self.deep_eval.get_model()
diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py
index 2363e29100..fe5e6ee668 100644
--- a/deepmd/pd/infer/deep_eval.py
+++ b/deepmd/pd/infer/deep_eval.py
@@ -506,6 +506,16 @@ def get_model_size(self) -> dict:
             "total": sum_param_des + sum_param_fit,
         }
 
+    def get_model(self):
+        """Get the Paddle model module (paddle.nn.Layer).
+
+        Returns
+        -------
+        paddle.nn.Layer
+            The Paddle model module.
+        """
+        return self.dp
+
     def eval_descriptor(
         self,
         coords: np.ndarray,
diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py
index 13bd4d2bf0..1e07b99261 100644
--- a/deepmd/pt/infer/deep_eval.py
+++ b/deepmd/pt/infer/deep_eval.py
@@ -706,6 +706,16 @@ def get_observed_types(self) -> dict:
             "observed_type": sort_element_type(observed_type_list),
         }
 
+    def get_model(self):
+        """Get the PyTorch model module (nn.Module).
+
+        Returns
+        -------
+        torch.nn.Module
+            The PyTorch model module.
+        """
+        return self.dp
+
     def eval_descriptor(
         self,
         coords: np.ndarray,
diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py
index a7682d2e58..d9260a23a0 100644
--- a/deepmd/tf/infer/deep_eval.py
+++ b/deepmd/tf/infer/deep_eval.py
@@ -1126,6 +1126,16 @@ def get_model_def_script(self) -> dict:
         model_def_script = script.decode("utf-8")
         return json.loads(model_def_script)["model"]
 
+    def get_model(self):
+        """Get the TensorFlow graph.
+
+        Returns
+        -------
+        tf.Graph
+            The TensorFlow graph.
+        """
+        return self.graph
+
 
 class DeepEvalOld:
     # old class for DipoleChargeModifier only
diff --git a/source/tests/infer/test_get_model.py b/source/tests/infer/test_get_model.py
new file mode 100644
index 0000000000..1307837320
--- /dev/null
+++ b/source/tests/infer/test_get_model.py
@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+from deepmd.infer.deep_eval import (
+    DeepEval,
+)
+
+from ..consistent.common import (
+    parameterized,
+)
+from .case import (
+    get_cases,
+)
+
+
+@parameterized(
+    (
+        "se_e2_a",
+        "fparam_aparam",
+    ),  # key
+    (".pb", ".pth"),  # model extension
+)
+class TestGetModelMethod(unittest.TestCase):
+    """Test the new get_model method functionality."""
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        key, extension = cls.param
+        cls.case = get_cases()[key]
+        cls.model_name = cls.case.get_model(extension)
+        cls.dp = DeepEval(cls.model_name)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cls.dp = None
+
+    def test_get_model_method_exists(self):
+        """Test that get_model method exists."""
+        self.assertTrue(
+            hasattr(self.dp, "get_model"), "DeepEval should have get_model method"
+        )
+
+    def test_get_model_returns_valid_object(self):
+        """Test that get_model returns a valid model object."""
+        model = self.dp.get_model()
+        self.assertIsNotNone(model, "get_model should return a non-None object")
+
+    def test_get_model_backend_specific(self):
+        """Test that get_model returns the expected type for each backend."""
+        key, extension = self.param
+        model = self.dp.get_model()
+
+        if extension == ".pth":
+            # For PyTorch models, should return nn.Module-like object
+            import torch
+
+            self.assertTrue(
+                hasattr(model, "forward") or isinstance(model, torch.nn.Module),
+                "PyTorch model should be or behave like nn.Module",
+            )
+            # Check if it has model attribute (ModelWrapper)
+            if hasattr(model, "model"):
+                self.assertTrue(
+                    hasattr(model.model, "__getitem__"),
+                    "PyTorch ModelWrapper should have model dict",
+                )
+        elif extension == ".pb":
+            # For TensorFlow models, should return graph
+            try:
+                # Should be a TensorFlow graph or have graph-like properties
+                self.assertTrue(
+                    hasattr(model, "get_operations")
+                    or str(type(model)).find("Graph") >= 0,
+                    "TensorFlow model should be a graph or graph-like object",
+                )
+            except ImportError:
+                # If TensorFlow not available, skip this assertion
+                pass
+
+    def test_get_model_consistency(self):
+        """Test that get_model always returns the same object."""
+        model1 = self.dp.get_model()
+        model2 = self.dp.get_model()
+        # Should return the same object (not necessarily equal, but same reference)
+        self.assertIs(
+            model1, model2, "get_model should return consistent object reference"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

From e5c79686e94b095019bb30f5607541102b27691b Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 28 Aug 2025 17:36:16 +0000
Subject: [PATCH 3/6] fix(infer): implement get_model for JAX backend and
 return BaseModel for all non-TF backends

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>
---
 deepmd/jax/infer/deep_eval.py | 10 ++++++++++
 deepmd/pd/infer/deep_eval.py  | 13 +++++++++----
 deepmd/pt/infer/deep_eval.py  | 28 ++++++++++++++++++++++++----
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py
index acfd42b66a..9f13929228 100644
--- a/deepmd/jax/infer/deep_eval.py
+++ b/deepmd/jax/infer/deep_eval.py
@@ -420,3 +420,13 @@ def _get_output_shape(self, odef, nframes, natoms):
     def get_model_def_script(self) -> dict:
         """Get model definition script."""
         return json.loads(self.dp.get_model_def_script())
+
+    def get_model(self):
+        """Get the JAX model as BaseModel.
+
+        Returns
+        -------
+        BaseModel
+            The JAX model as BaseModel instance.
+        """
+        return self.dp
diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py
index fe5e6ee668..fa0a7f6281 100644
--- a/deepmd/pd/infer/deep_eval.py
+++ b/deepmd/pd/infer/deep_eval.py
@@ -11,6 +11,9 @@
 import paddle
 
 from deepmd.dpmodel.common import PRECISION_DICT as NP_PRECISION_DICT
+from deepmd.dpmodel.model.base_model import (
+    BaseModel,
+)
 from deepmd.dpmodel.output_def import (
     ModelOutputDef,
     OutputVariableCategory,
@@ -507,14 +510,16 @@ def get_model_size(self) -> dict:
         }
 
     def get_model(self):
-        """Get the Paddle model module (paddle.nn.Layer).
+        """Get the Paddle model as BaseModel.
 
         Returns
         -------
-        paddle.nn.Layer
-            The Paddle model module.
+        BaseModel
+            The Paddle model converted to BaseModel.
         """
-        return self.dp
+        # Convert Paddle model to BaseModel by serializing and deserializing
+        model_dict = self.dp.model["Default"].serialize()
+        return BaseModel.deserialize(model_dict)
 
     def eval_descriptor(
         self,
diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py
index 1e07b99261..d4f98a5111 100644
--- a/deepmd/pt/infer/deep_eval.py
+++ b/deepmd/pt/infer/deep_eval.py
@@ -13,6 +13,9 @@
 import torch
 
 from deepmd.dpmodel.common import PRECISION_DICT as NP_PRECISION_DICT
+from deepmd.dpmodel.model.base_model import (
+    BaseModel,
+)
 from deepmd.dpmodel.output_def import (
     ModelOutputDef,
     OutputVariableCategory,
@@ -707,14 +710,31 @@ def get_observed_types(self) -> dict:
         }
 
     def get_model(self):
-        """Get the PyTorch model module (nn.Module).
+        """Get the PyTorch model as BaseModel.
 
         Returns
         -------
-        torch.nn.Module
-            The PyTorch model module.
+        BaseModel
+            The PyTorch model converted to BaseModel.
         """
-        return self.dp
+        # Convert PyTorch model to BaseModel by serializing and deserializing
+        if str(self.model_path).endswith(".pth"):
+            # For JIT models (.pth), we need to reconstruct the original model first
+            from deepmd.pt.model.model import (
+                get_model,
+            )
+
+            # The JIT model should have model_def_script
+            model_def_script = self.model_def_script
+            model = get_model(model_def_script)
+            # Load state dict with strict=False to handle compression info differences
+            model.load_state_dict(self.dp.model["Default"].state_dict(), strict=False)
+            model_dict = model.serialize()
+        else:
+            # For regular models (.pt), we can serialize directly
+            model_dict = self.dp.model["Default"].serialize()
+
+        return BaseModel.deserialize(model_dict)
 
     def eval_descriptor(
         self,

From 9eff90a07ade936cabac26a7199019f46c356d62 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 28 Aug 2025 18:15:40 +0000
Subject: [PATCH 4/6] fix(infer): return backend-specific models instead of
 converting to dpmodel BaseModel

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>
---
 deepmd/pd/infer/deep_eval.py         | 11 +++--------
 deepmd/pt/infer/deep_eval.py         | 26 +++-----------------------
 source/tests/infer/test_get_model.py | 24 ++++++++++++++----------
 3 files changed, 20 insertions(+), 41 deletions(-)

diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py
index fa0a7f6281..374ec03366 100644
--- a/deepmd/pd/infer/deep_eval.py
+++ b/deepmd/pd/infer/deep_eval.py
@@ -11,9 +11,6 @@
 import paddle
 
 from deepmd.dpmodel.common import PRECISION_DICT as NP_PRECISION_DICT
-from deepmd.dpmodel.model.base_model import (
-    BaseModel,
-)
 from deepmd.dpmodel.output_def import (
     ModelOutputDef,
     OutputVariableCategory,
@@ -510,16 +507,14 @@ def get_model_size(self) -> dict:
         }
 
     def get_model(self):
-        """Get the Paddle model as BaseModel.
+        """Get the Paddle model.
 
         Returns
         -------
         BaseModel
-            The Paddle model converted to BaseModel.
+            The Paddle model instance.
         """
-        # Convert Paddle model to BaseModel by serializing and deserializing
-        model_dict = self.dp.model["Default"].serialize()
-        return BaseModel.deserialize(model_dict)
+        return self.dp.model["Default"]
 
     def eval_descriptor(
         self,
diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py
index d4f98a5111..68cf2845fb 100644
--- a/deepmd/pt/infer/deep_eval.py
+++ b/deepmd/pt/infer/deep_eval.py
@@ -13,9 +13,6 @@
 import torch
 
 from deepmd.dpmodel.common import PRECISION_DICT as NP_PRECISION_DICT
-from deepmd.dpmodel.model.base_model import (
-    BaseModel,
-)
 from deepmd.dpmodel.output_def import (
     ModelOutputDef,
     OutputVariableCategory,
@@ -710,31 +707,14 @@ def get_observed_types(self) -> dict:
         }
 
     def get_model(self):
-        """Get the PyTorch model as BaseModel.
+        """Get the PyTorch model.
 
         Returns
         -------
         BaseModel
-            The PyTorch model converted to BaseModel.
+            The PyTorch model instance.
         """
-        # Convert PyTorch model to BaseModel by serializing and deserializing
-        if str(self.model_path).endswith(".pth"):
-            # For JIT models (.pth), we need to reconstruct the original model first
-            from deepmd.pt.model.model import (
-                get_model,
-            )
-
-            # The JIT model should have model_def_script
-            model_def_script = self.model_def_script
-            model = get_model(model_def_script)
-            # Load state dict with strict=False to handle compression info differences
-            model.load_state_dict(self.dp.model["Default"].state_dict(), strict=False)
-            model_dict = model.serialize()
-        else:
-            # For regular models (.pt), we can serialize directly
-            model_dict = self.dp.model["Default"].serialize()
-
-        return BaseModel.deserialize(model_dict)
+        return self.dp.model["Default"]
 
     def eval_descriptor(
         self,
diff --git a/source/tests/infer/test_get_model.py b/source/tests/infer/test_get_model.py
index 1307837320..1cd447c585 100644
--- a/source/tests/infer/test_get_model.py
+++ b/source/tests/infer/test_get_model.py
@@ -51,19 +51,23 @@ def test_get_model_backend_specific(self):
         model = self.dp.get_model()
 
         if extension == ".pth":
-            # For PyTorch models, should return nn.Module-like object
-            import torch
+            # For PyTorch models, should return the PyTorch BaseModel
+            from deepmd.pt.model.model.model import BaseModel as PTBaseModel
 
+            self.assertIsInstance(
+                model,
+                PTBaseModel,
+                "PyTorch model should return PyTorch BaseModel instance",
+            )
+            # Check if it has common model methods
             self.assertTrue(
-                hasattr(model, "forward") or isinstance(model, torch.nn.Module),
-                "PyTorch model should be or behave like nn.Module",
+                hasattr(model, "get_type_map"),
+                "PyTorch BaseModel should have get_type_map method",
+            )
+            self.assertTrue(
+                hasattr(model, "get_rcut"),
+                "PyTorch BaseModel should have get_rcut method",
             )
-            # Check if it has model attribute (ModelWrapper)
-            if hasattr(model, "model"):
-                self.assertTrue(
-                    hasattr(model.model, "__getitem__"),
-                    "PyTorch ModelWrapper should have model dict",
-                )
         elif extension == ".pb":
             # For TensorFlow models, should return graph
             try:

From 233e0eaf3a660ac20e33c40b4e26f9ec8347e09e Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 29 Aug 2025 02:04:02 +0000
Subject: [PATCH 5/6] fix(test): update PyTorch .pth model test to expect
 TorchScript ScriptModule

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>
---
 source/tests/infer/test_get_model.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/source/tests/infer/test_get_model.py b/source/tests/infer/test_get_model.py
index 1cd447c585..4c52dda0a1 100644
--- a/source/tests/infer/test_get_model.py
+++ b/source/tests/infer/test_get_model.py
@@ -51,22 +51,28 @@ def test_get_model_backend_specific(self):
         model = self.dp.get_model()
 
         if extension == ".pth":
-            # For PyTorch models, should return the PyTorch BaseModel
-            from deepmd.pt.model.model.model import BaseModel as PTBaseModel
+            # For PyTorch .pth models (TorchScript), should return torch.jit.ScriptModule
+            import torch
 
             self.assertIsInstance(
                 model,
-                PTBaseModel,
-                "PyTorch model should return PyTorch BaseModel instance",
+                torch.jit.ScriptModule,
+                "PyTorch .pth model should return TorchScript ScriptModule instance",
+            )
+            # TorchScript modules are also nn.Module instances
+            self.assertIsInstance(
+                model,
+                torch.nn.Module,
+                "PyTorch .pth model should be a torch.nn.Module instance",
             )
             # Check if it has common model methods
             self.assertTrue(
                 hasattr(model, "get_type_map"),
-                "PyTorch BaseModel should have get_type_map method",
+                "PyTorch model should have get_type_map method",
             )
             self.assertTrue(
                 hasattr(model, "get_rcut"),
-                "PyTorch BaseModel should have get_rcut method",
+                "PyTorch model should have get_rcut method",
             )
         elif extension == ".pb":
             # For TensorFlow models, should return graph

From 7ba9fad93f82ed8e50aa0e097d6d8a421bdbb495 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 2 Sep 2025 09:11:00 +0000
Subject: [PATCH 6/6] Addressing PR comments

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>
---
 deepmd/dpmodel/infer/deep_eval.py | 2 +-
 deepmd/infer/deep_eval.py         | 4 ++--
 deepmd/jax/infer/deep_eval.py     | 2 +-
 deepmd/pd/infer/deep_eval.py      | 6 +++++-
 deepmd/pt/infer/deep_eval.py      | 6 +++++-
 deepmd/tf/infer/deep_eval.py      | 2 +-
 6 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py
index e19da62103..4b30d97c29 100644
--- a/deepmd/dpmodel/infer/deep_eval.py
+++ b/deepmd/dpmodel/infer/deep_eval.py
@@ -393,7 +393,7 @@ def get_model_def_script(self) -> dict:
         """Get model definition script."""
         return json.loads(self.dp.get_model_def_script())
 
-    def get_model(self):
+    def get_model(self) -> "BaseModel":
         """Get the dpmodel BaseModel.
 
         Returns
diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py
index e97f1cb0a6..49c9576afd 100644
--- a/deepmd/infer/deep_eval.py
+++ b/deepmd/infer/deep_eval.py
@@ -342,7 +342,7 @@ def get_observed_types(self) -> dict:
         raise NotImplementedError("Not implemented in this backend.")
 
     @abstractmethod
-    def get_model(self):
+    def get_model(self) -> Any:
         """Get the model module implemented by the deep learning framework.
 
         For PyTorch, this returns the nn.Module. For Paddle, this returns
@@ -700,7 +700,7 @@ def get_observed_types(self) -> dict:
         """Get observed types (elements) of the model during data statistics."""
         return self.deep_eval.get_observed_types()
 
-    def get_model(self):
+    def get_model(self) -> Any:
         """Get the model module implemented by the deep learning framework.
 
         For PyTorch, this returns the nn.Module. For Paddle, this returns
diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py
index 9f13929228..2e74c15fff 100644
--- a/deepmd/jax/infer/deep_eval.py
+++ b/deepmd/jax/infer/deep_eval.py
@@ -421,7 +421,7 @@ def get_model_def_script(self) -> dict:
         """Get model definition script."""
         return json.loads(self.dp.get_model_def_script())
 
-    def get_model(self):
+    def get_model(self) -> Any:
         """Get the JAX model as BaseModel.
 
         Returns
diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py
index 374ec03366..61c3f9e9a3 100644
--- a/deepmd/pd/infer/deep_eval.py
+++ b/deepmd/pd/infer/deep_eval.py
@@ -46,6 +46,10 @@
 if TYPE_CHECKING:
     import ase.neighborlist
 
+    from deepmd.pd.model.model.model import (
+        BaseModel,
+    )
+
 
 class DeepEval(DeepEvalBackend):
     """Paddle backend implementation of DeepEval.
@@ -506,7 +510,7 @@ def get_model_size(self) -> dict:
             "total": sum_param_des + sum_param_fit,
         }
 
-    def get_model(self):
+    def get_model(self) -> "BaseModel":
         """Get the Paddle model.
 
         Returns
diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py
index 68cf2845fb..25caf12b64 100644
--- a/deepmd/pt/infer/deep_eval.py
+++ b/deepmd/pt/infer/deep_eval.py
@@ -75,6 +75,10 @@
 if TYPE_CHECKING:
     import ase.neighborlist
 
+    from deepmd.pt.model.model.model import (
+        BaseModel,
+    )
+
 log = logging.getLogger(__name__)
 
 
@@ -706,7 +710,7 @@ def get_observed_types(self) -> dict:
             "observed_type": sort_element_type(observed_type_list),
         }
 
-    def get_model(self):
+    def get_model(self) -> "BaseModel":
         """Get the PyTorch model.
 
         Returns
diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py
index d9260a23a0..75440accb9 100644
--- a/deepmd/tf/infer/deep_eval.py
+++ b/deepmd/tf/infer/deep_eval.py
@@ -1126,7 +1126,7 @@ def get_model_def_script(self) -> dict:
         model_def_script = script.decode("utf-8")
         return json.loads(model_def_script)["model"]
 
-    def get_model(self):
+    def get_model(self) -> "tf.Graph":
         """Get the TensorFlow graph.
 
         Returns