jax在下面pycharm中有提示你在哪里下载,所以如果下次遇到有什么要下载的包都可以尝试这种方式去下载(在找找pycharm有没有提示,然后直接点网站看教程,这嘎是官网的教程,很详细)
?
?为了下载一个版本只有JAX的cpu,可能要用到当地的开发,你可以用.....(我英语很差的!)
依次输入以下代码
pip install --upgrade pip
pip install --upgrade "jax[cpu]"
pip install --upgrade pip
# Installs the wheel compatible with CUDA 11 and cuDNN 8.2 or newer.
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html # Note: wheels only available on linux.
你可以为jaxlib显式指定一个特定的CUDA和cuDNN版本: ?
pip install --upgrade pip
# Installs the wheel compatible with Cuda 11 and cudnn 8.2 or newer.
pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_releases.html
# Installs the wheel compatible with Cuda 11 and cudnn 8.0.5 or newer.
pip install jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_releases.html
你可以通过命令找到你的CUDA版本: ?
nvcc --version
就可以了,还有什么不会的下面提问
|