diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index 98185c0117..a5dcfb16dd 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -276,9 +276,9 @@ def call( gg = self.cal_g(tr, tt) gg = np.mean(gg, axis=2) # nf x nloc x ng x 1 - xyz_scatter += gg + xyz_scatter += gg * (self.sel[tt] / self.nnei) - res_rescale = 1.0 / 10.0 + res_rescale = 1.0 / 5.0 res = xyz_scatter * res_rescale res = res.reshape(nf, nloc, -1).astype(GLOBAL_NP_FLOAT_PRECISION) return res, None, None, None, ww diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index 1debcc8caf..16721fbe5e 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -258,9 +258,9 @@ def forward( # nfnl x nt x ng gg = ll.forward(ss) gg = torch.mean(gg, dim=1).unsqueeze(1) - xyz_scatter += gg + xyz_scatter += gg * (self.sel[ii] / self.nnei) - res_rescale = 1.0 / 10.0 + res_rescale = 1.0 / 5.0 result = xyz_scatter * res_rescale result = result.view(-1, nloc, self.filter_neuron[-1]) return ( diff --git a/source/tests/consistent/descriptor/test_se_e2_a.py b/source/tests/consistent/descriptor/test_se_e2_a.py index 0243a77044..b8f4205d09 100644 --- a/source/tests/consistent/descriptor/test_se_e2_a.py +++ b/source/tests/consistent/descriptor/test_se_e2_a.py @@ -51,7 +51,7 @@ def data(self) -> dict: precision, ) = self.param return { - "sel": [10, 10], + "sel": [9, 10], "rcut_smth": 5.80, "rcut": 6.00, "neuron": [6, 12, 24], diff --git a/source/tests/consistent/descriptor/test_se_r.py b/source/tests/consistent/descriptor/test_se_r.py index 354ae1cc99..8b835f3b5c 100644 --- a/source/tests/consistent/descriptor/test_se_r.py +++ b/source/tests/consistent/descriptor/test_se_r.py @@ -51,7 +51,7 @@ def data(self) -> dict: precision, ) = self.param return { - "sel": [10, 10], + "sel": [9, 10], "rcut_smth": 5.80, "rcut": 6.00, "neuron": [6, 12, 24],