Skip to content

Commit

Permalink
Update repformers.py
Browse files Browse the repository at this point in the history
  • Loading branch information
iProzd committed May 9, 2024
1 parent d1e38ad commit 9d0ad7f
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion deepmd/pt/model/descriptor/repformers.py
Original file line number Diff line number Diff line change
Expand Up @@ -400,7 +400,7 @@ def forward(
nframes, nloc, nnei = nlist.shape
nall = extended_coord.view(nframes, -1).shape[1] // 3
atype = extended_atype[:, :nloc]
# nf x nloc x nnei
# nb x nloc x nnei
exclude_mask = self.emask(nlist, extended_atype)
nlist = nlist * exclude_mask
# nb x nloc x nnei x 4, nb x nloc x nnei x 3, nb x nloc x nnei x 1
Expand Down

0 comments on commit 9d0ad7f

Please sign in to comment.