【fixed】No GPU/TPU found, falling back to CPU.或者XlaRuntimeErrorFAILED_PRECONDITION: DNN library initi

问题

用jax报错No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

或者
XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed.

解决

pip install --upgrade pip

# CUDA 12 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# CUDA 11 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html