do not change diff directly

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Oct 25, 2024 · d0b576f · d0b576f
1 parent 0517b59
commit d0b576f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py
@@ -62,8 +62,8 @@ def _make_env_mat(
     diff = coord_r - coord_l
     # nf x nloc x nnei
     # the grad of JAX vector_norm is NaN at x=0
-    diff = xp.where(xp.abs(diff) < 1e-30, xp.full_like(diff, 1e-30), diff)
-    length = xp.linalg.vector_norm(diff, axis=-1, keepdims=True)
+    diff_ = xp.where(xp.abs(diff) < 1e-30, xp.full_like(diff, 1e-30), diff)
+    length = xp.linalg.vector_norm(diff_, axis=-1, keepdims=True)
     # for index 0 nloc atom
     length = length + xp.astype(~xp.expand_dims(mask, axis=-1), length.dtype)
     t0 = 1 / (length + protection)