set the tpu cores in the env file

nitins17 · nitins17 · commit 5b450af2223f · 2024-09-20T17:07:14.000Z
diff --git a/ci/envs/run_tests/pytest_tpu b/ci/envs/run_tests/pytest_tpu
@@ -19,8 +19,10 @@ export JAXCI_INSTALL_WHEELS_LOCALLY=1
 
 export JAXCI_RUN_PYTEST_TPU=1
 
+export JAXCI_TPU_CORES=8
+
 # Disable x64 mode
 export JAX_ENABLE_X64=0
 
 # Clone XLA at HEAD.
-export JAXCI_CLONE_MAIN_XLA=1
+export JAXCI_CLONE_MAIN_XLA=1
