diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 0a995eaf94..6124356e4e 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -88,10 +88,9 @@ def compute_output_stats( rcond: Optional[float] = None, atom_ener: Optional[List[float]] = None, model_forward: Optional[Callable[..., torch.Tensor]] = None, - ): key_mapping = { - "polar":"polarizability", + "polar": "polarizability", "energy": "energy", "dos": "dos", "dipole": "dipole", @@ -288,7 +287,6 @@ def compute_output_stats_with_atomic( which will be subtracted from the energy label of the data. The difference will then be used to calculate the delta complement energy bias for each type. """ - atomic_label_name, global_label_name = "atom_" + key, key if stat_file_path is not None: diff --git a/source/tests/pt/model/test_polar_stat.py b/source/tests/pt/model/test_polar_stat.py index d3f6334182..052163b60c 100644 --- a/source/tests/pt/model/test_polar_stat.py +++ b/source/tests/pt/model/test_polar_stat.py @@ -62,9 +62,9 @@ def test_atomic_consistency(self): def test_global_consistency(self): self.sampled[0]["find_atom_polarizability"] = -1 - self.sampled[0]["polarizability"] = self.sampled[0][ - "atom_polarizability" - ].sum(dim=1) + self.sampled[0]["polarizability"] = self.sampled[0]["atom_polarizability"].sum( + dim=1 + ) self.all_stat["find_atom_polarizability"] = [-1] self.all_stat["polarizability"] = [ self.all_stat["atom_polarizability"][0].sum(axis=1)