From 85f4e14b0ff9b64a808d777b606653557e9233ce Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 18 Apr 2024 01:09:31 -0400 Subject: [PATCH 1/2] fix(tf): fix foat32 for exclude_types in se_atten_v2 Fix type issue in previous PR #3651. Signed-off-by: Jinzhe Zeng --- deepmd/tf/descriptor/se_atten.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 3ca763870b..bde8365775 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -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: From 005b5ed979f23258c5329a231c5d766aeecbb71f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 18 Apr 2024 05:10:45 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/descriptor/se_atten.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index bde8365775..4018ebab14 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -705,7 +705,11 @@ def _pass_filter( ), ) self.recovered_switch *= tf.reshape( - tf.slice(tf.reshape(tf.cast(mask, self.filter_precision), [-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: