I just installed CUDA on my Ubuntu 19.10 desktop, and it was a pretty quick process, if non-obivous.
- Installed base CUDA with
apt install nvidia-cuda-toolkit - Downloaded libcudnn from nvidia and installed with
dpkg -i. (I couldn't find it in the official Ubuntu repo, let me know if it is there) - Set
XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/lib/cuda/so JAX can find the installed libraries.
Tags: misc