fix input in test_network.py

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Nov 13, 2024 · 36b087d · 36b087d
1 parent e43134d
commit 36b087d
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/source/tests/common/dpmodel/test_network.py b/source/tests/common/dpmodel/test_network.py
@@ -8,6 +8,9 @@
 
 import numpy as np
 
+from deepmd.dpmodel.common import (
+    get_xp_precision,
+)
 from deepmd.dpmodel.utils import (
     EmbeddingNet,
     FittingNet,
@@ -46,7 +49,9 @@ def test_serialize_deserize(self):
             inp_shap = [ni]
             if ashp is not None:
                 inp_shap = ashp + inp_shap
-            inp = np.arange(np.prod(inp_shap)).reshape(inp_shap)
+            inp = np.arange(
+                np.prod(inp_shap), dtype=get_xp_precision(np, prec)
+            ).reshape(inp_shap)
             np.testing.assert_allclose(nl0.call(inp), nl1.call(inp))
 
     def test_shape_error(self):
@@ -168,7 +173,7 @@ def test_embedding_net(self):
                 resnet_dt=idt,
             )
             en1 = EmbeddingNet.deserialize(en0.serialize())
-            inp = np.ones([ni])
+            inp = np.ones([ni], dtype=get_xp_precision(np, prec))
             np.testing.assert_allclose(en0.call(inp), en1.call(inp))
 
 
@@ -191,7 +196,7 @@ def test_fitting_net(self):
                 bias_out=bo,
             )
             en1 = FittingNet.deserialize(en0.serialize())
-            inp = np.ones([ni])
+            inp = np.ones([ni], dtype=get_xp_precision(np, prec))
             en0.call(inp)
             en1.call(inp)
             np.testing.assert_allclose(en0.call(inp), en1.call(inp))