find element in list rather than tuple, which easily makes bugs

deepmd/model_format/fitting.py

@@ -200,29 +200,29 @@ def output_def(self):
         )
 
     def __setitem__(self, key, value):
-        if key in ("bias_atom_e"):
+        if key in ["bias_atom_e"]:
             self.bias_atom_e = value
-        elif key in ("fparam_avg"):
+        elif key in ["fparam_avg"]:
             self.fparam_avg = value
-        elif key in ("fparam_inv_std"):
+        elif key in ["fparam_inv_std"]:
             self.fparam_inv_std = value
-        elif key in ("aparam_avg"):
+        elif key in ["aparam_avg"]:
             self.aparam_avg = value
-        elif key in ("aparam_inv_std"):
+        elif key in ["aparam_inv_std"]:
             self.aparam_inv_std = value
         else:
             raise KeyError(key)
 
     def __getitem__(self, key):
-        if key in ("bias_atom_e"):
+        if key in ["bias_atom_e"]:
             return self.bias_atom_e
-        elif key in ("fparam_avg"):
+        elif key in ["fparam_avg"]:
             return self.fparam_avg
-        elif key in ("fparam_inv_std"):
+        elif key in ["fparam_inv_std"]:
             return self.fparam_inv_std
-        elif key in ("aparam_avg"):
+        elif key in ["aparam_avg"]:
             return self.aparam_avg
-        elif key in ("aparam_inv_std"):
+        elif key in ["aparam_inv_std"]:
             return self.aparam_inv_std
         else:
             raise KeyError(key)

deepmd/pt/model/task/ener.py

@@ -166,30 +166,30 @@ def output_def(self) -> FittingOutputDef:
         )
 
     def __setitem__(self, key, value):
-        if key in ("bias_atom_e"):
+        if key in ["bias_atom_e"]:
             # correct bias_atom_e shape. user may provide stupid  shape
             self.bias_atom_e = value
-        elif key in ("fparam_avg"):
+        elif key in ["fparam_avg"]:
             self.fparam_avg = value
-        elif key in ("fparam_inv_std"):
+        elif key in ["fparam_inv_std"]:
             self.fparam_inv_std = value
-        elif key in ("aparam_avg"):
+        elif key in ["aparam_avg"]:
             self.aparam_avg = value
-        elif key in ("aparam_inv_std"):
+        elif key in ["aparam_inv_std"]:
             self.aparam_inv_std = value
         else:
             raise KeyError(key)
 
     def __getitem__(self, key):
-        if key in ("bias_atom_e"):
+        if key in ["bias_atom_e"]:
             return self.bias_atom_e
-        elif key in ("fparam_avg"):
+        elif key in ["fparam_avg"]:
             return self.fparam_avg
-        elif key in ("fparam_inv_std"):
+        elif key in ["fparam_inv_std"]:
             return self.fparam_inv_std
-        elif key in ("aparam_avg"):
+        elif key in ["aparam_avg"]:
             return self.aparam_avg
-        elif key in ("aparam_inv_std"):
+        elif key in ["aparam_inv_std"]:
             return self.aparam_inv_std
         else:
             raise KeyError(key)

source/tests/common/test_model_format_utils.py

@@ -471,3 +471,22 @@ def test_self_exception(
                 with self.assertRaises(ValueError) as context:
                     ret0 = ifn0(dd[0], atype, fparam=ifp, aparam=iap)
                     self.assertIn("input aparam", context.exception)
+
+    def test_get_set(self):
+        ifn0 = InvarFitting(
+            "energy",
+            self.nt,
+            3,
+            1,
+        )
+        rng = np.random.default_rng()
+        foo = rng.normal([3, 4])
+        for ii in [
+            "bias_atom_e",
+            "fparam_avg",
+            "fparam_inv_std",
+            "aparam_avg",
+            "aparam_inv_std",
+        ]:
+            ifn0[ii] = foo
+            np.testing.assert_allclose(foo, ifn0[ii])

source/tests/pt/test_ener_fitting.py

@@ -160,3 +160,22 @@ def test_jit(
                 use_tebd=(not distinguish_types),
             ).to(env.DEVICE)
             torch.jit.script(ft0)
+
+    def test_get_set(self):
+        ifn0 = InvarFitting(
+            "energy",
+            self.nt,
+            3,
+            1,
+        )
+        rng = np.random.default_rng()
+        foo = rng.normal([3, 4])
+        for ii in [
+            "bias_atom_e",
+            "fparam_avg",
+            "fparam_inv_std",
+            "aparam_avg",
+            "aparam_inv_std",
+        ]:
+            ifn0[ii] = torch.tensor(foo, dtype=dtype, device=env.DEVICE)
+            np.testing.assert_allclose(foo, ifn0[ii].detach().cpu().numpy())