RuntimeError: Unknown backend gpu

if you want to run on GPU then the first thing you would do is to check your GPU driver and CUDA version. For nvidia for example I run nvidia-smi to check that. Depending on that information you should install the following list (assuming you have CUDA==12.2) pip install nvidia-cuda-cupti-cu12==12.2.131 pip install nvidia-cuda-nvcc-cu12==12.2.140 pip install nvidia-cuda-nvrtc-cu12==12.2.140 pip install nvidia-cuda-runtime-cu12==12.2.140 pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html This did it for me.. You will then need to make sure you have enough GPU memory. I found it extremely difficult to port this code on aws GPU images. Maybe you will have a better luck with your system..

If you run the code on CPU then these two worked for me.. pip install jax==0.4.23 pip install jaxlib==0.4.23

Hope this helps..

+---------------------------------------------------------------------------------------+ | Processes: | | GPU GI CI PID Type Process name GPU Memory | | ID ID Usage | |=======================================================================================| | No running processes found | +---------------------------------------------------------------------------------------+

ecmwf-lab / ai-models-graphcast

RuntimeError: Unknown backend gpu #16