CUDA backend failed to initialize

[mio-19]

I got the warning messsage on the "Load Model" stage. This might prevent the notebook to use gpu

WARNING:jax._src.xla_bridge:CUDA backend failed to initialize: Unable to load CUDA. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

i got the same warning... have you solved the problem now?

[sandywu]

replace jax[cuda11_local] wth jax[cuda12_pip]==0.4.20