Correct libtpu_storage_path for stable libtpu builds (#8576)

pytorch · Jan 15, 2025 · 79de480 · 79de480
1 parent e513e13
commit 79de480
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -70,14 +70,16 @@
 _jax_version = f'0.4.39'
 _jaxlib_version = f'0.4.39'
 _libtpu_wheel_name = f'libtpu-{_libtpu_version}'
+_libtpu_storage_directory = 'libtpu-lts-releases'
 
 if USE_NIGHTLY:
   _libtpu_version += f".dev{_date}"
   _jax_version += f".dev{_date}"
   _jaxlib_version += f".dev{_date}"
   _libtpu_wheel_name += f".dev{_date}+nightly"
+  _libtpu_storage_directory = 'libtpu-nightly-releases'
 
-_libtpu_storage_path = f'https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu/{_libtpu_wheel_name}-py3-none-linux_x86_64.whl'
+_libtpu_storage_path = f'https://storage.googleapis.com/{_libtpu_storage_directory}/wheels/libtpu/{_libtpu_wheel_name}-py3-none-linux_x86_64.whl'
 
 
 def _get_build_mode():