Update CUDA Version [Bug]

[hideosnes]

The Colab (and probably also the Huggingface Space) exits with a CUDA mismatch notification (Jax requires a newer version of Cuda). What I could figure out is that CUDAv is a transitive dependency of Jax and Nvidia's own pip packages don't have the necessary constrain as of yet.

Got it to run on a fork with !pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html and by adding the KeyArray for Jax manually. However, that's a wonky hack that might cease to work (very) soon.

[hideosnes]

Got it to work: !apt-get -y install cuda-toolkit-12-4 should do the trick! Hopefully Nvidia will update their pip packages any time soon, then regular !pip should work in the colab again.