Fix: ZBL state_dict duplicated keys (#3432)

deepmodeling · Mar 12, 2024 · bc35ac9 · bc35ac9
1 parent a88a213
commit bc35ac9
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
@@ -361,8 +361,6 @@ def __init__(
         models = [dp_model, zbl_model]
         super().__init__(models, type_map, **kwargs)
         self.model_def_script = ""
-        self.dp_model = dp_model
-        self.zbl_model = zbl_model
 
         self.sw_rmin = sw_rmin
         self.sw_rmax = sw_rmax
@@ -391,8 +389,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
@@ -406,7 +404,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,