diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index 637eec6496..e8560edadd 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -56,8 +56,12 @@ def change_output_def(self) -> None: ) ) self.deep_eval.output_def = self.output_def - self.deep_eval._OUTDEF_DP2BACKEND[self.get_property_name()] = f"atom_{self.get_property_name()}" - self.deep_eval._OUTDEF_DP2BACKEND[f"{self.get_property_name()}_redu"] = self.get_property_name() + self.deep_eval._OUTDEF_DP2BACKEND[self.get_property_name()] = ( + f"atom_{self.get_property_name()}" + ) + self.deep_eval._OUTDEF_DP2BACKEND[f"{self.get_property_name()}_redu"] = ( + self.get_property_name() + ) @property def task_dim(self) -> int: @@ -125,7 +129,9 @@ def eval( atomic_property = results[self.get_property_name()].reshape( nframes, natoms, self.get_task_dim() ) - property = results[f"{self.get_property_name()}_redu"].reshape(nframes, self.get_task_dim()) + property = results[f"{self.get_property_name()}_redu"].reshape( + nframes, self.get_task_dim() + ) if atomic: return ( diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index 713aee6f5c..8ae1e7a264 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/deepmd/pt/model/model/property_model.py @@ -62,8 +62,12 @@ def forward( do_atomic_virial=do_atomic_virial, ) model_predict = {} - model_predict[f"atom_{self.get_property_name()}"] = model_ret[self.get_property_name()] - model_predict[self.get_property_name()] = model_ret[f"{self.get_property_name()}_redu"] + model_predict[f"atom_{self.get_property_name()}"] = model_ret[ + self.get_property_name() + ] + model_predict[self.get_property_name()] = model_ret[ + f"{self.get_property_name()}_redu" + ] if "mask" in model_ret: model_predict["mask"] = model_ret["mask"] return model_predict @@ -107,8 +111,12 @@ def forward_lower( extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) model_predict = {} - model_predict[f"atom_{self.get_property_name()}"] = model_ret[self.get_property_name()] - model_predict[self.get_property_name()] = model_ret[f"{self.get_property_name()}_redu"] + model_predict[f"atom_{self.get_property_name()}"] = model_ret[ + self.get_property_name() + ] + model_predict[self.get_property_name()] = model_ret[ + f"{self.get_property_name()}_redu" + ] if "mask" in model_ret: model_predict["mask"] = model_ret["mask"] return model_predict