fix(tf): fix foat32 for exclude_types in se_atten_v2 (#3682)

Fix type issue in previous PR #3651. --------- Signed-off-by: Jinzhe Zeng <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> (cherry picked from commit cac8715)
deepmodeling · Jul 3, 2024 · 93ac49b · 93ac49b
1 parent 87dfdf3
commit 93ac49b
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py
@@ -673,7 +673,11 @@ 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: