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

Add version constraints to CUDA pip wheel dependencies.

Browse Source 
Fixes https://github.com/google/jax/issues/15267
This commit is contained in:
Peter Hawkins 2023-03-28 21:55:32 -04:00
parent 8c4fed6410
commit 705b5cc000
1 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
setup.py
View File

@ -96,13 +96,13 @@ setup(
'cuda11_pip': [
f"jaxlib=={_current_jaxlib_version}+cuda11.cudnn{_default_cuda11_cudnn_version}",
"nvidia-cublas-cu11",
"nvidia-cuda-nvcc-cu11",
"nvidia-cuda-runtime-cu11",
"nvidia-cudnn-cu11",
"nvidia-cufft-cu11",
"nvidia-cusolver-cu11",
"nvidia-cusparse-cu11",
"nvidia-cublas-cu11>=11.11",
"nvidia-cuda-nvcc-cu11>=11.8",
"nvidia-cuda-runtime-cu11>=11.8",
"nvidia-cudnn-cu11>=8.6",
"nvidia-cufft-cu11>=10.9",
"nvidia-cusolver-cu11>=11.4",
"nvidia-cusparse-cu11>=11.7",
],
'cuda12_pip': [