fix(2024Q1): build nlist faster with torch.amax (deepmodeling#3827)

iProzd · May 28, 2024 · 22026a8 · 22026a8
2 parents eed7c8a + a64a806
commit 22026a8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py
@@ -326,7 +326,7 @@ def extend_coord_with_ghosts(
         # +1: central cell
         nbuff = torch.ceil(rcut / to_face).to(torch.long)
         # 3
-        nbuff = torch.max(nbuff, dim=0, keepdim=False).values
+        nbuff = torch.amax(nbuff, dim=0)  # faster than torch.max
         nbuff_cpu = nbuff.cpu()
         xi = torch.arange(-nbuff_cpu[0], nbuff_cpu[0] + 1, 1, device="cpu")
         yi = torch.arange(-nbuff_cpu[1], nbuff_cpu[1] + 1, 1, device="cpu")