diff --git a/deepmd/dpmodel/descriptor/se_atten_v2.py b/deepmd/dpmodel/descriptor/se_atten_v2.py index c247ccc573..8044e10eca 100644 --- a/deepmd/dpmodel/descriptor/se_atten_v2.py +++ b/deepmd/dpmodel/descriptor/se_atten_v2.py @@ -165,7 +165,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": embeddings = data.pop("embeddings") type_embedding = data.pop("type_embedding") attention_layers = data.pop("attention_layers") - env_mat = data.pop("env_mat") + data.pop("env_mat") embeddings_strip = data.pop("embeddings_strip") obj = cls(**data) diff --git a/source/tests/pt/model/test_se_atten_v2.py b/source/tests/pt/model/test_se_atten_v2.py index 9afe330048..5d68b2f440 100644 --- a/source/tests/pt/model/test_se_atten_v2.py +++ b/source/tests/pt/model/test_se_atten_v2.py @@ -105,7 +105,7 @@ def test_jit( self, ): rng = np.random.default_rng() - nf, nloc, nnei = self.nlist.shape + _, _, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) dstd = rng.normal(size=(self.nt, nnei, 4)) dstd = 0.1 + np.abs(dstd) @@ -123,8 +123,6 @@ def test_jit( [False, True], # use_econf_tebd ): dtype = PRECISION_DICT[prec] - rtol, atol = get_tols(prec) - err_msg = f"idt={idt} prec={prec}" # dpa1 new impl dd0 = DescrptSeAttenV2( self.rcut, @@ -140,6 +138,4 @@ def test_jit( ) dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) dd0.se_atten.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) - # dd1 = DescrptDPA1.deserialize(dd0.serialize()) - model = torch.jit.script(dd0) - # model = torch.jit.script(dd1) + _ = torch.jit.script(dd0) \ No newline at end of file