Detecting CUDA/cuDNN

[r-wedeen]

Hi,

First, I want to thank you for creating this repository -- it looks extremely useful.

I am currently trying to run turbozero on a LambdaLabs GPU instance which runs CUDA12. I followed your instructions and ran the commands

• pip install poetry==1.7.1
• poetry install
• poetry source add jax https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
• poetry add jax[cuda12_pip]==0.4.24

However, when I run

• poetry shell
• python
• import jax
• print(jax.devices())

I get:

CUDA backend failed to initialize: Unable to load cuDNN. Is it installed? (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.) [CpuDevice(id=0)]

My apologies in advance if this is a silly/simple error or if I should ask this somewhere else -- I am new to poetry, CUDA, etc. and any pointers or comments to help me along would be greatly appreciated. Thank you!

[lowrollr]

Hi, thanks for trying out my project.

Possible the setup guide needs updating, I can take a look very soon.