Update debug-metadata.yml

huggingface · Dec 4, 2024 · bd377bf · bd377bf
1 parent 653ddaa
commit bd377bf
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/.github/workflows/debug-metadata.yml b/.github/workflows/debug-metadata.yml
@@ -24,5 +24,14 @@ jobs:
             --shm-size=16G \
             -v /mnt/hf_cache:/mnt/cache/ \
             -e PJRT_DEVICE=TPU \
+            -e TPU_TOPOLOGY \
+            -e TPU_WORKER_ID \
+            -e TPU_SKIP_MDS_QUERY \
+            -e TPU_TOPOLOGY_WRAP \
+            -e TPU_CHIPS_PER_HOST_BOUNDS \
+            -e TPU_ACCELERATOR_TYPE \
+            -e TPU_TOPOLOGY_ALT \
+            -e TPU_HOST_BOUNDS \
+            -e TPU_WORKER_HOSTNAMES \
             us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla@sha256:8f1dcd5b03f993e4da5c20d17c77aff6a5f22d5455f8eb042d2e4b16ac460526 \
-            bash -c "python -c 'import torch_xla.core.xla_model as xm; assert xm.xla_device().type == \"xla\", \"XLA device not available\"'"
+            bash -c "python -c 'import torch_xla.core.xla_model as xm; assert xm.xla_device().type == \"xla\", \"XLA device not available\"'"