diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 81f2cc8cf0..1a883a50a2 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -62,6 +62,7 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, distinguish_types: bool = False, + rcond: Optional[float] = None, **kwargs, ): """Construct a fitting net for energy. @@ -87,6 +88,7 @@ def __init__( self.activation_function = activation_function self.precision = precision self.prec = PRECISION_DICT[self.precision] + self.rcond = rcond if bias_atom_e is None: bias_atom_e = np.zeros([self.ntypes, self.dim_out]) bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) @@ -217,8 +219,7 @@ def compute_output_stats(self, merged): input_natoms = [item["real_natoms_vec"] for item in merged] else: input_natoms = [item["natoms"] for item in merged] - tmp = compute_output_bias(energy, input_natoms) - bias_atom_e = tmp[:, 0] + bias_atom_e = compute_output_bias(energy, input_natoms, rcond=self.rcond) return {"bias_atom_e": bias_atom_e} def init_fitting_stat(self, bias_atom_e=None, **kwargs): @@ -244,6 +245,7 @@ def serialize(self) -> dict: "precision": self.precision, "distinguish_types": self.distinguish_types, "nets": self.filter_layers.serialize(), + "rcond": self.rcond, "@variables": { "bias_atom_e": to_numpy_array(self.bias_atom_e), "fparam_avg": to_numpy_array(self.fparam_avg), @@ -259,7 +261,6 @@ def serialize(self) -> dict: # "use_aparam_as_mask": self.use_aparam_as_mask , # "spin": self.spin , ## NOTICE: not supported by far - "rcond": None, "tot_ener_zero": False, "trainable": True, "atom_ener": None,