ci(docker): update torch-xla to 2.4.0 (#89)

Now it has been released, it makes sense to use that version.
huggingface · Sep 3, 2024 · 5eb6cf3 · 5eb6cf3
1 parent d4e2294
commit 5eb6cf3
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/.github/workflows/test-pytorch-xla-tpu-tgi-nightly.yml b/.github/workflows/test-pytorch-xla-tpu-tgi-nightly.yml
@@ -17,7 +17,7 @@ jobs:
     runs-on: optimum-tpu
     container:
       # Use a nightly image that works with TPU (release was not working)
-      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla@sha256:48b1d3ab080613fd88234019daf77ef7812b518acb13c54ddad03bf770d6ac57
+      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_tpuvm
       options: --shm-size "16gb" --ipc host --privileged
     env:
       PJRT_DEVICE: TPU

diff --git a/.github/workflows/test-pytorch-xla-tpu-tgi.yml b/.github/workflows/test-pytorch-xla-tpu-tgi.yml
@@ -20,7 +20,7 @@ jobs:
     runs-on: optimum-tpu
     container:
       # Use a nightly image that works with TPU (release was not working)
-      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla@sha256:48b1d3ab080613fd88234019daf77ef7812b518acb13c54ddad03bf770d6ac57
+      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_tpuvm
       options: --shm-size "16gb" --ipc host --privileged
     env:
       PJRT_DEVICE: TPU

diff --git a/.github/workflows/test-pytorch-xla-tpu.yml b/.github/workflows/test-pytorch-xla-tpu.yml
@@ -20,7 +20,7 @@ jobs:
     runs-on: optimum-tpu
     container:
       # Use a nightly image that works with TPU (release was not working)
-      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla@sha256:48b1d3ab080613fd88234019daf77ef7812b518acb13c54ddad03bf770d6ac57
+      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_tpuvm
       options: --shm-size "16gb" --ipc host --privileged
     env:
       PJRT_DEVICE: TPU