From 6b80175de547f40a201a583c1bdf62b3aa7e06ea Mon Sep 17 00:00:00 2001
From: cenzhiyao <2523403608@qq.com>
Date: Wed, 18 Dec 2024 19:58:31 +0800
Subject: [PATCH] add torch nightly build in run-gh200-test.sh

---
 .buildkite/run-gh200-test.sh | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/.buildkite/run-gh200-test.sh b/.buildkite/run-gh200-test.sh
index 6cf071b02b23a..ff8ae37742dfb 100644
--- 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