Merge branch 'devel' into fix-model-def-script

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Mar 12, 2024 · 7094868 · 7094868
2 parents 72fbcfd + bc35ac9
commit 7094868
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -360,8 +360,6 @@ def __init__(
     ):
         models = [dp_model, zbl_model]
         super().__init__(models, type_map, **kwargs)
-        self.dp_model = dp_model
-        self.zbl_model = zbl_model
 
         self.sw_rmin = sw_rmin
         self.sw_rmax = sw_rmax
@@ -390,8 +388,8 @@ def compute_or_load_stat(
         stat_file_path
             The dictionary of paths to the statistics files.
         """
-        self.dp_model.compute_or_load_stat(sampled_func, stat_file_path)
-        self.zbl_model.compute_or_load_stat(sampled_func, stat_file_path)
+        self.models[0].compute_or_load_stat(sampled_func, stat_file_path)
+        self.models[1].compute_or_load_stat(sampled_func, stat_file_path)
 
     def change_energy_bias(self):
         # need to implement
@@ -405,7 +403,7 @@ def serialize(self) -> dict:
                 "@version": 1,
                 "type": "zbl",
                 "models": LinearEnergyAtomicModel.serialize(
-                    [self.dp_model, self.zbl_model], self.type_map
+                    [self.models[0], self.models[1]], self.type_map
                 ),
                 "sw_rmin": self.sw_rmin,
                 "sw_rmax": self.sw_rmax,