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
就可以了,还有什么不会的下面提问
欢迎分享,转载请注明来源:内存溢出
评论列表(0条)