From 9d0ad7fae0bf426f032b6edd251581bd0f99f5b2 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Thu, 9 May 2024 15:43:46 +0800
Subject: [PATCH] Update repformers.py

---
 deepmd/pt/model/descriptor/repformers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py
index b2df49c964..e352c6b40c 100644
--- a/deepmd/pt/model/descriptor/repformers.py
+++ b/deepmd/pt/model/descriptor/repformers.py
@@ -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