From b64101b87db61a278b3ca9efe1f967d1d5a79ab3 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Tue, 27 Feb 2024 15:14:03 +0000
Subject: [PATCH] WIP

---
 .github/workflows/test_tpu_tgi.yml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.github/workflows/test_tpu_tgi.yml b/.github/workflows/test_tpu_tgi.yml
index 869eb770..afd6a907 100644
--- a/.github/workflows/test_tpu_tgi.yml
+++ b/.github/workflows/test_tpu_tgi.yml
@@ -19,6 +19,8 @@ jobs:
   do-the-job:
     name: Run TGI tests
     runs-on: [optimum-tpu] # run the configured Google TPU
+    env:
+      PJRT_DEVICE: TPU
     steps:
       - name: Checkout
         uses: actions/checkout@v2
@@ -28,6 +30,7 @@ jobs:
           source optimum-tpu-pytorch/bin/activate
           python -m pip install -U pip
           python -m pip install "torch~=2.2.0" "torch_xla[tpu]~=2.2.0" -f https://storage.googleapis.com/libtpu-releases/index.html
+          python -c "import torch_xla.core.xla_model as xm; assert xm.xla_device().type == 'xla', 'XLA device not available'"
       - name: Run TGI server python tests
         run: |
           source optimum-tpu-pytorch/bin/activate