fix(infer): return backend-specific models instead of converting to dpmodel BaseModel

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

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,

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,

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: