numagic / lumos

scalable accelerated optimal control
MIT License
16 stars 0 forks source link

update jax installation to use new gpu release #86

Closed yunlongxu-numagic closed 1 year ago

yunlongxu-numagic commented 1 year ago

the gpu release has changed, we need to update the doc, and the colab notebook

pip install "jax[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html