Update deepmd/dpmodel/fitting/general_fitting.py

Co-authored-by: Han Wang <[email protected]> Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Nov 13, 2024 · 30d33b7 · 30d33b7
1 parent 41d80b9
commit 30d33b7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py
@@ -457,4 +457,4 @@ def _call_common(
         exclude_mask = self.emask.build_type_exclude_mask(atype)
         # nf x nloc x nod
         outs = xp.where(exclude_mask[:, :, None], outs, xp.zeros_like(outs))
-        return {self.var_name: xp.astype(outs, descriptor.dtype)}
+        return {self.var_name: outs)}