Add support for cuda 11.3
VanWieren opened this issue · 2 comments
VanWieren commented
As far as I can tell the only gpu support is for cuda <=11.2
hawkinsp commented
Actually: good news! The CUDA 11.1 wheel should work for 11.2 and 11.3. We're going to document this in the next release of jaxlib
.
With the released CUDA 11.1 jaxlib wheel, you may need to set the environment variable XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/local/cuda/
, but that should be fixed with the next release.
Please try it out and let us know if it works!
hawkinsp commented
We just released jaxlib 0.1.66, which should support CUDA 11.3 without any workarounds. Install the cuda111
variant of the wheel, which should work on CUDA 11.3 without any special setup.
Hope that helps!