From c2a669539f194dec07c025932b7e8a8f15df5028 Mon Sep 17 00:00:00 2001
From: Manfei <41607353+ManfeiBai@users.noreply.github.com>
Date: Thu, 17 Aug 2023 18:25:03 -0700
Subject: [PATCH] Update llama2-model.libsonnet

---
 tests/pytorch/nightly/llama2-model.libsonnet | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/pytorch/nightly/llama2-model.libsonnet b/tests/pytorch/nightly/llama2-model.libsonnet
index b4a372ecc..bbe129ab2 100644
--- a/tests/pytorch/nightly/llama2-model.libsonnet
+++ b/tests/pytorch/nightly/llama2-model.libsonnet
@@ -151,8 +151,11 @@ local utils = import 'templates/utils.libsonnet';
         sudo apt update
         sudo apt-get -y install libopenblas-dev
         pip3 uninstall -y torch torch_xla
-        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly-cp38-cp38-linux_x86_64.whl
-        pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp38-cp38-linux_x86_64.whl
+        ## pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly-cp38-cp38-linux_x86_64.whl
+        ## pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp38-cp38-linux_x86_64.whl
+        pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch-nightly+20230422-cp38-cp38-linux_x86_64.whl
+        pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/cuda/118/torch_xla-nightly+20230422-cp38-cp38-linux_x86_64.whl
+        pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/cuda/118/torchvision-nightly+20230422-cp38-cp38-linux_x86_64.whl
         pip3 uninstall -y libtpu-nightly
         pip3 install torch_xla[tpuvm] --user
         # show current path