add torch nightly build in run-gh200-test.sh

vllm-project · Dec 18, 2024 · 6b80175 · 6b80175
1 parent 6bddcd4
commit 6b80175
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/.buildkite/run-gh200-test.sh b/.buildkite/run-gh200-test.sh
@@ -15,6 +15,11 @@ apt-get update -y \
   && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
   && python3 --version && python3 -m pip --version
 
+# Install the nightly version of torch and torchvision
+python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cu124 "torch==2.6.0.dev20241210+cu124"
+python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cu124 "torchvision==0.22.0.dev20241215"
+
+# Skip the new torch installation during build since we are using the specified version
 python3 use_existing_torch.py
 
 # Try building the docker image