diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py
index a4f81a23a5..1060b40ce1 100644
--- a/deepmd/pt/utils/nlist.py
+++ b/deepmd/pt/utils/nlist.py
@@ -457,7 +457,7 @@ def extend_coord_with_ghosts(
         xyz = xyz.view(-1, 3)
         xyz = xyz.to(device=device, non_blocking=True)
         # ns x 3
-        shift_idx = xyz[torch.argsort(torch.norm(xyz, dim=1))]
+        shift_idx = xyz[torch.argsort(torch.linalg.norm(xyz, dim=-1))]
         ns, _ = shift_idx.shape
         nall = ns * nloc
         # nf x ns x 3