Skip to content

Commit

Permalink
fix(tf): fix foat32 for exclude_types in se_atten_v2
Browse files Browse the repository at this point in the history
Fix type issue in previous PR deepmodeling#3651.

Signed-off-by: Jinzhe Zeng <[email protected]>
  • Loading branch information
njzjz authored Apr 18, 2024
1 parent a04b46d commit 85f4e14
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion deepmd/tf/descriptor/se_atten.py
Original file line number Diff line number Diff line change
Expand Up @@ -705,7 +705,7 @@ def _pass_filter(
),
)
self.recovered_switch *= tf.reshape(
tf.slice(tf.reshape(mask, [-1, 4]), [0, 0], [-1, 1]),
tf.slice(tf.reshape(tf.cast(mask, self.filter_precision), [-1, 4]), [0, 0], [-1, 1]),
[-1, natoms[0], self.sel_all_a[0]],
)
else:
Expand Down

0 comments on commit 85f4e14

Please sign in to comment.