fix(pt-expt): align spin metadata output definitions

Author: OutisLi

deepmd/pt_expt/infer/deep_eval.py

@@ -143,21 +143,12 @@ def _init_from_model_json(self, model_json_str: str) -> None:
         self._sel = list(self._dpmodel.get_sel())
         self._mixed_types = bool(self._dpmodel.mixed_types())
         if self._is_spin:
-            self._model_output_def = ModelOutputDef(
-                FittingOutputDef(
-                    [
-                        OutputVariableDef(
-                            "energy",
-                            shape=[1],
-                            reducible=True,
-                            r_differentiable=True,
-                            c_differentiable=True,
-                            atomic=True,
-                            magnetic=True,
-                        )
-                    ]
-                )
-            )
+            spin_fitting_defs = self._dpmodel.model_output_def().def_outp.get_data()
+            # Keep only physical fitting outputs; mask is derived by ModelOutputDef.
+            fitting_defs = [
+                vdef for name, vdef in spin_fitting_defs.items() if name != "mask"
+            ]
+            self._model_output_def = ModelOutputDef(FittingOutputDef(fitting_defs))
         else:
             self._model_output_def = ModelOutputDef(self._dpmodel.atomic_output_def())
 

deepmd/pt_expt/utils/serialization.py

@@ -255,9 +255,15 @@ def _collect_metadata(model: torch.nn.Module, is_spin: bool = False) -> dict:
     The ``fitting_output_defs`` list is also included so that
     ``ModelOutputDef`` can be reconstructed without loading the full model.
     """
-    fitting_output_def = model.atomic_output_def()
+    if is_spin:
+        fitting_output_def = model.model_output_def().def_outp
+    else:
+        fitting_output_def = model.atomic_output_def()
     fitting_output_defs = []
     for vdef in fitting_output_def.get_data().values():
+        # Keep metadata aligned with physical fitting outputs only.
+        if is_spin and vdef.name == "mask":
+            continue
         fitting_output_defs.append(
             {
                 "name": vdef.name,