set dtype

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Nov 1, 2024 · 41863c9 · 41863c9
1 parent 1b3ea6b
commit 41863c9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/dpmodel/utils/safe_gradient.py b/deepmd/dpmodel/utils/safe_gradient.py
@@ -28,5 +28,5 @@ def safe_for_vector_norm(x, /, *, axis=None, keepdims=False, ord=2):
         xp.linalg.vector_norm(
             xp.where(mask, x, xp.ones_like(x)), axis=axis, keepdims=keepdims, ord=ord
         ),
-        xp.zeros_like(mask_squeezed),
+        xp.zeros_like(mask_squeezed, dtype=x.dtype),
     )