Skip to content

Commit 8c6e2ea

Browse files
authored
fix(tf): always use float64 for the global tensor (deepmodeling#4735)
Fix deepmodeling#4734. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved the calculation of global output by applying an additional transformation to per-atom outputs before aggregation, resulting in more accurate global predictions. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
1 parent f0b1069 commit 8c6e2ea

1 file changed

Lines changed: 2 additions & 1 deletion

File tree

deepmd/tf/model/tensor.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
from deepmd.tf.env import (
88
MODEL_VERSION,
9+
global_cvt_2_ener_float,
910
tf,
1011
)
1112
from deepmd.tf.utils.type_embed import (
@@ -173,7 +174,7 @@ def build(
173174
if "global" not in self.model_type:
174175
gname = "global_" + self.model_type
175176
atom_out = tf.reshape(output, [-1, natomsel, nout])
176-
global_out = tf.reduce_sum(atom_out, axis=1)
177+
global_out = tf.reduce_sum(global_cvt_2_ener_float(atom_out), axis=1)
177178
global_out = tf.reshape(global_out, [-1, nout], name="o_" + gname + suffix)
178179

179180
out_cpnts = tf.split(atom_out, nout, axis=-1)

0 commit comments

Comments
 (0)