fix(pt): resolve issues with dipole charge modifier

- Use pickle instead of json for modifier serialization to handle np.ndarray
- Fix device placement in EwaldReal class by explicitly setting device and dtype
- Add proper error handling in test case to ensure torch default dtype is restored

Author: ChiahsinChu

deepmd/pt/entrypoints/main.py

@@ -404,6 +404,8 @@ def freeze(
     extra_files = {"modifier_data": ""}
     dm = tester.modifier
     if dm is not None:
+        # dict from dm.serialize() includes np.ndarray
+        # use pickle rather than json
         bytes_data = pickle.dumps(dm.serialize())
         extra_files = {"modifier_data": bytes_data}
     torch.jit.save(

deepmd/pt/modifier/dipole_charge.py

@@ -85,7 +85,7 @@ def __init__(
             rspace=False,
             kappa=ewald_beta,
             spacing=ewald_h,
-        ).to(env.GLOBAL_PT_FLOAT_PRECISION)
+        ).to(device=env.DEVICE, dtype=env.GLOBAL_PT_FLOAT_PRECISION)
         self.er = torch.jit.script(er)
         self.er.eval()
         self.placeholder_pairs = torch.ones((1, 2), device=env.DEVICE, dtype=torch.long)

source/tests/pt/modifier/test_dipole_charge.py

@@ -160,30 +160,31 @@ def test_consistency(self):
         dtype = torch.get_default_dtype()
         torch.set_default_dtype(torch.float64)
 
-        coord, box, atype = ref_data()
+        try:
+            coord, box, atype = ref_data()
 
-        pt_data = self.dm_pt.eval_np(
-            coord=coord,
-            atype=atype,
-            box=box,
-        )
-        tf_data = self.dm_tf.eval(
-            coord=coord,
-            box=box,
-            atype=atype.reshape(-1),
-        )
-        tol = 1e-6
-        output_names = ["energy", "force", "virial"]
-        for ii, name in enumerate(output_names):
-            np.testing.assert_allclose(
-                pt_data[ii].reshape(-1),
-                tf_data[ii].reshape(-1),
-                atol=tol,
-                rtol=tol,
-                err_msg=f"Mismatch in {name}",
+            pt_data = self.dm_pt.eval_np(
+                coord=coord,
+                atype=atype,
+                box=box,
             )
-
-        torch.set_default_dtype(dtype)
+            tf_data = self.dm_tf.eval(
+                coord=coord,
+                box=box,
+                atype=atype.reshape(-1),
+            )
+            tol = 1e-6
+            output_names = ["energy", "force", "virial"]
+            for ii, name in enumerate(output_names):
+                np.testing.assert_allclose(
+                    pt_data[ii].reshape(-1),
+                    tf_data[ii].reshape(-1),
+                    atol=tol,
+                    rtol=tol,
+                    err_msg=f"Mismatch in {name}",
+                )
+        finally:
+            torch.set_default_dtype(dtype)
 
     def test_serialize(self):
         """Test the serialize method of DipoleChargeModifier."""