mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Set JAX_PLATFORMS=tpu,cpu on TPUs

Browse Source
This commit is contained in:
Ran Ran 2022-10-03 22:52:03 +00:00
parent 015a12e155
commit 78a7e161bb
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jax/_src/cloud_tpu_init.py
View File

@ -48,6 +48,7 @@ def cloud_tpu_init():
libtpu.configure_library_path()
os.environ.setdefault('GRPC_VERBOSITY', 'ERROR')
os.environ.setdefault('JAX_PLATFORMS', 'tpu,cpu')
os.environ['TPU_ML_PLATFORM'] = 'JAX'
# If the user has set any topology-related env vars, don't set any