WARNING - No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

[zohimchandani]

I am pretty sure that the box I am connecting to via ssh has 20 CPU cores and 2 GPU cores however the GPUs are not recognized by Jax. Any ideas on how to solve this? Thanks.

[YouJiacheng]

Do you install jaxlib with GPU support?
try to run:

pip install -U jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Or

pip install -U jaxlib==0.3.14+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

[Xiang-Pan]

Same issue, does not work

[skr3178]

Do you install jaxlib with GPU support? try to run:

pip install -U jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Or

pip install -U jaxlib==0.3.14+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

First reference worked.
https://github.com/lucidrains/progen

[danizil]

Hi I was stuck on this issue for several hours and got the answer that solved everything for me:
instead of conda installing cudatoolkit, do:
conda install cudatoolkit-dev=11.3.1 -c conda-forge
(might not require that specific version, should probably be >= to your cuda version) apparently cudatoolkit lacks the compiler nvcc that JAX needs.
The thing is that on the server I work on, the nvcc compiler was not installed for each user separately but the command nvcc -V did work because it was installed for everyone.
In summary:
install cudatoolkit-dev and cuDNN with anaconda.
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html (will also install jaxlib)
Hope this helps!

[YouJiacheng]

conda install cuda -c nvidia can install the newest version of nvcc.(and a lots of tools)

[wzh-miasanmia]

I had this problem too, but my cause was because of a version conflict between jax and jaxlib, which seemed to automatically become the cpu version when updating jax;
My solution was:

pip uninstall jax jaxlib

pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

change the second command to your cuda version if you are not 12.x
hope this works :)