Update deepmd/pt/model/descriptor/se_atten.py

Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Signed-off-by: Anyang Peng <137014849+anyangml@users.noreply.github.com>
deepmodeling · May 9, 2024 · 53446dd · 53446dd
1 parent 1e7fd2f
commit 53446dd
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py
@@ -464,8 +464,7 @@ def forward(
             protection=self.env_protection,
         )
         nlist_mask = nlist != -1
-        nlist_copy = nlist.detach().clone()
-        nlist_copy[nlist == -1] = 0
+        nlist = torch.where(nlist == -1, 0, nlist)
         sw = torch.squeeze(sw, -1)
         # beyond the cutoff sw should be 0.0
         sw = sw.masked_fill(~nlist_mask, 0.0)