From 7214cd46b6adbe2018824f1064dd4318e79e5e2b Mon Sep 17 00:00:00 2001
From: Pauline Bailly-Masson <155966238+paulinebm@users.noreply.github.com>
Date: Wed, 17 Apr 2024 09:39:22 +0200
Subject: [PATCH] Update test-pytorch-xla-tpu-tgi.yml

---
 .github/workflows/test-pytorch-xla-tpu-tgi.yml | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/test-pytorch-xla-tpu-tgi.yml b/.github/workflows/test-pytorch-xla-tpu-tgi.yml
index 3d5c7f79..dda06167 100644
--- a/.github/workflows/test-pytorch-xla-tpu-tgi.yml
+++ b/.github/workflows/test-pytorch-xla-tpu-tgi.yml
@@ -20,14 +20,13 @@ jobs:
     runs-on: tpu-runner-dind2
     container:
       # Use an image that works with TPU with Pytorch 2.3.0 (release was not working)
-      image: python:3.10
-      #image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla@sha256:8f1dcd5b03f993e4da5c20d17c77aff6a5f22d5455f8eb042d2e4b16ac460526
+      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla@sha256:8f1dc5b03f993e4da5c20d17c77aff6a5f22d5455f8eb042d2e4b16ac460526
       options: --shm-size "16gb" --ipc host --privileged
     env:
       PJRT_DEVICE: TPU
     steps:
-      #- name: Checkout
-        #uses: actions/checkout@v4
+      - name: Checkout
+        uses: actions/checkout@v4
       - name: slepp
         run: |
           apt update