From c027ae677c3f5b8f8ce1fc26104f8eb89373095d Mon Sep 17 00:00:00 2001 From: Hank Han Date: Mon, 11 May 2026 16:22:20 +0800 Subject: [PATCH] ci: align torch CUDA wheel checks (#24940) --- scripts/ci/cuda/ci_install_dependency.sh | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index 5d8e0cd69..398ba3993 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -277,7 +277,20 @@ install_sglang_kernel() { # TODO: Remove after torch 2.11 where cu13 is enabled by default TORCH_CUDA_VER=$(python3 -c "import torch; v=torch.version.cuda; parts=v.split('.'); print(f'cu{parts[0]}{parts[1]}')") echo "Detected torch CUDA version: ${TORCH_CUDA_VER}" + TORCHAUDIO_CUDA_VER=$(pip show torchaudio 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed -n 's/.*+\(cu[0-9][0-9]*\)$/\1/p' || true) + TORCHVISION_CUDA_VER=$(pip show torchvision 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed -n 's/.*+\(cu[0-9][0-9]*\)$/\1/p' || true) + REINSTALL_TORCH=false if [ "${TORCH_CUDA_VER}" != "${CU_VERSION}" ]; then + REINSTALL_TORCH=true + else + for cuda_ver in "${TORCHAUDIO_CUDA_VER}" "${TORCHVISION_CUDA_VER}"; do + if [ -n "${cuda_ver}" ] && [ "${cuda_ver}" != "${CU_VERSION}" ]; then + REINSTALL_TORCH=true + break + fi + done + fi + if [ "${REINSTALL_TORCH}" = true ]; then TORCH_VER=$(pip show torch 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//') TORCHAUDIO_VER=$(pip show torchaudio 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//') TORCHVISION_VER=$(pip show torchvision 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')