Jax cannot access GPU when generating predictions

google-deepmind / graphcast

Apache License 2.0

4.54k stars 572 forks source link

Jax cannot access GPU when generating predictions #27

Open scploeger opened 10 months ago

scploeger commented 10 months ago

When generating predictions, this error comes up:

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

Using an A100 environment on colab with the stock notebook except with the modifications mentioned here.

tvonich commented 8 months ago

Any luck on this? I'm having a similar issue.

AndrewY-bot commented 5 months ago

Maybe is the version of cuda and cudnn not match the jaxlib version? You can find right version here. By the way, you need to uninstall jax&jaxlib of old version first.