Merge branch 'split-model-atomic-model' of https://github.com/njzjz/d…

…eepmd-kit into split-model-atomic-model
deepmodeling · Mar 9, 2024 · ed831c8 · ed831c8
2 parents 8f6614c + e22bc64
commit ed831c8
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py
@@ -26,6 +26,7 @@
     BaseModel,
 )
 from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
     ModelOutputDef,
     OutputVariableCategory,
     OutputVariableOperation,
@@ -470,4 +471,8 @@ def mixed_types(self) -> bool:
             """
             return self.atomic_model.mixed_types()
 
+        def atomic_output_def(self) -> FittingOutputDef:
+            """Get the output def of the atomic model."""
+            return self.atomic_model.atomic_output_def()
+
     return CM