From 6879de25b8e2e221ea13053690f1a89776f9664b Mon Sep 17 00:00:00 2001
From: Manfei <41607353+ManfeiBai@users.noreply.github.com>
Date: Mon, 9 Sep 2024 18:07:04 -0700
Subject: [PATCH] Update common.libsonnet to use built nightly wheel

---
 dags/legacy_test/tests/pytorch/nightly/common.libsonnet | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/dags/legacy_test/tests/pytorch/nightly/common.libsonnet b/dags/legacy_test/tests/pytorch/nightly/common.libsonnet
index 15b588e4..1b14eba3 100644
--- a/dags/legacy_test/tests/pytorch/nightly/common.libsonnet
+++ b/dags/legacy_test/tests/pytorch/nightly/common.libsonnet
@@ -96,10 +96,12 @@ local volumes = import 'templates/volumes.libsonnet';
         sudo apt install -y libopenblas-base
         # for huggingface tests
         sudo apt install -y libsndfile-dev
-        pip3 install --user --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
-        pip install --user \
+        pip3 install --user \
+          https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly-cp310-cp310-linux_x86_64.whl
+        pip3 install --user \
           'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp310-cp310-linux_x86_64.whl' \
           -f https://storage.googleapis.com/libtpu-releases/index.html
+        pip3 install --user --pre --no-deps torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
         pip3 install pillow
         git clone --depth=1 https://github.com/pytorch/pytorch.git
         cd pytorch