From 51e5f492dcfdffe3301d64e1698d453abb440ae7 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Tue, 27 Feb 2024 01:24:32 +0800
Subject: [PATCH] use torch.logical_and

---
 deepmd/pt/utils/env_mat_stat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py
index 08365ad7c9..7e235f3762 100644
--- a/deepmd/pt/utils/env_mat_stat.py
+++ b/deepmd/pt/utils/env_mat_stat.py
@@ -146,7 +146,7 @@ def iter(
                 ).view(-1, 1),
             )
             # shape: (ntypes, nloc, nnei)
-            type_idx = (type_idx.unsqueeze(-1) * exclude_mask) == 1
+            type_idx = torch.logical_and(type_idx.unsqueeze(-1), exclude_mask)
             for type_i in range(self.descriptor.get_ntypes()):
                 dd = env_mat[type_idx[type_i]]
                 dd = dd.reshape([-1, 4])  # typen_atoms * unmasked_nnei, 4