diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py
index 423f26c227..0fa51f5b83 100644
--- 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
index a29a2fbcff..b8b59f48ac 100644
--- 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(