Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <suvorovv@google.com>

Author: mydatascience

.github/workflows/tpu-tests.yml

@@ -42,6 +42,9 @@ jobs:
   tpu_unit_tests:
     needs: prelim
     runs-on: [linux-x86-ct5lp-224-8tpu]
+    container:
+      image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest
+      options: --privileged
     steps:
     - name: Checkout code
       uses: actions/checkout@v4
@@ -118,7 +121,10 @@ jobs:
 
   tpu_integration_tests:
     needs: prelim
-    runs-on: [self-hosted, linux-x86-ct5lp-224-8tpu]
+    runs-on: [linux-x86-ct5lp-224-8tpu]
+    container:
+      image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest
+      options: --privileged
     steps:
     - name: Checkout code
       uses: actions/checkout@v4