fix(pt): build nlist faster with torch.amax#3826

iProzd merged 1 commit intodeepmodeling:develdeepmodeling/deepmd-kit:develfrom

iProzd:fix_nlist_maxiProzd/deepmd-kit:fix_nlist_maxCopy head branch name to clipboard