dtype

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Mar 2, 2024 · 9e55925 · 9e55925
1 parent 80337f7
commit 9e55925
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py
@@ -24,7 +24,7 @@ def compute_smooth_weight(
         uu = (distance - rmin) / (rmax - rmin)
     with np.errstate(invalid="ignore"):
         vv = uu * uu * uu * (-6.0 * uu * uu + 15.0 * uu - 10.0) + 1.0
-    return np.where(mid_mask, vv, min_mask)
+    return np.where(mid_mask, vv, min_mask.astype(distance.dtype))
 
 
 def _make_env_mat(

diff --git a/deepmd/pt/utils/preprocess.py b/deepmd/pt/utils/preprocess.py
@@ -233,7 +233,7 @@ def compute_smooth_weight(distance, rmin: float, rmax: float):
     mid_mask = torch.logical_not(torch.logical_or(min_mask, max_mask))
     uu = (distance - rmin) / (rmax - rmin)
     vv = uu * uu * uu * (-6 * uu * uu + 15 * uu - 10) + 1
-    return torch.where(mid_mask, vv, min_mask)
+    return torch.where(mid_mask, vv, min_mask.to(dtype=distance.dtype))
 
 
 def make_env_mat(