use torch.logical_and

iProzd · Feb 26, 2024 · 51e5f49 · 51e5f49
1 parent 4606ab0
commit 51e5f49
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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