JAX 0.3.10: GPU not found?

[jecampagne]

Hello,

I just run an fresh JAX 0.3.10 version
jax 0.3.10
jaxlib 0.3.10

on a GPU I just get the following

 WARNING:absl:No GPU/TPU found, falling back to CPU.

while the same code on 0.3.5 was detecting the GPU ?! an idea?
Thanks

[jakevdp]

Can you say more about how you installed jaxlib?

[jecampagne]

Grrr.. Sorry I have forgotten the cuda update..

pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

Now it is ok. Sorry...