diff --git a/deepmd/pt/model/descriptor/repflows.py b/deepmd/pt/model/descriptor/repflows.py
index a9d2326b93..031287a797 100644
--- a/deepmd/pt/model/descriptor/repflows.py
+++ b/deepmd/pt/model/descriptor/repflows.py
@@ -496,8 +496,8 @@ def forward(
         a_sw = a_sw.masked_fill(~a_nlist_mask, 0.0)
         # set all padding positions to index of 0
         # if the a neighbor is real or not is indicated by nlist_mask
-        nlist[nlist == -1] = 0
-        a_nlist[a_nlist == -1] = 0
+        nlist = torch.where(nlist == -1, 0, nlist)
+        a_nlist = torch.where(a_nlist == -1, 0, a_nlist)
 
         # get node embedding
         # [nframes, nloc, tebd_dim]
diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py
index 75e2f97576..ca8da0ac1b 100644
--- a/deepmd/pt/model/descriptor/repformers.py
+++ b/deepmd/pt/model/descriptor/repformers.py
@@ -457,7 +457,7 @@ def forward(
 
         # set all padding positions to index of 0
         # if the a neighbor is real or not is indicated by nlist_mask
-        nlist[nlist == -1] = 0
+        nlist = torch.where(nlist == -1, 0, nlist)
         # nb x nall x ng1
         if comm_dict is None:
             assert mapping is not None