From bc35ac99800fe4a6bd4609c33a12cd024e1469c2 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 12 Mar 2024 09:30:16 +0800 Subject: [PATCH] Fix: ZBL state_dict duplicated keys (#3432) --- deepmd/pt/model/atomic_model/linear_atomic_model.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 3fb3ee90dd..1649921270 100644 --- 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,