No GPU/TPU found error

[sunhoro]

Any help would be appreciated.
I am running on two PC with Windows 11 and Windows 10, each with Nvidia GPU, and facing same issue on both.

All installation was done correctly but keep getting "No GPU/TPU found, falling back to CPU" error which still runs fine but super slow.
I think cuda and jax might not be compatible at this point but not sure.

I made sure everything is up to date including pip, cuda, and jax but still the same issue.

• cuda = 12.2
• jax = 0.4.13
• pip = 23.1.2

[Sk-bishnoi]

same error getting, if you got @sunhoro solution please let me know
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
output/1.png
final output/1.png

[njustlxk]

You need to install the GPU version of JAX, which is often not compatible with Windows