CUDA 11.0 and cuDNN / NVIDIA driver versions

Hi,

The Docker image suggests that the CUDA version in the GPU runtime would be 11.0, which has known issues and should be deprecated. This also limits the use of key packages like JAX which requires at least CUDA 11.1.

Is there any chance the CUDA version can be bumped to >= 11.1? Would it also be possible to use a nvidia/cuda image with known cuDNN versions, such as nvidia/cuda/11.5.1-cudnn8-runtime-ubuntu20.04 in case it helps with env setup? Would we also know NVIDIA driver version in the GPU runtime?

Thanks!