#激活虚拟环境#
在此之前确定好自己电脑的CUDA版本以及CUDNN版本
按照我自己电脑的版本配置,我先后输入
pip install jax[cuda12_pip]==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
安装完之后检查
python
import jax.numpy as np
中途没有报错就安装成功啦!