Setup Docker with GPU support on Ubuntu

setup-docker-with-GPU-support.sh

# STEP 1
$ sudo apt update && \
sudo apt -qq -y install \
    gcc \
    make \
    software-properties-common \
    apt-transport-https \
    ca-certificates \
    curl \
    git \
    jq && \
BASE_URL=https://us.download.nvidia.com/tesla && \
DRIVER_VERSION=525.105.17 && \
curl -fSsl -O $BASE_URL/$DRIVER_VERSION/NVIDIA-Linux-x86_64-$DRIVER_VERSION.run && \
sudo sh NVIDIA-Linux-x86_64-$DRIVER_VERSION.run && \
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add - && \
sudo add-apt-repository \
   "deb [arch=amd64] https://download.docker.com/linux/ubuntu \
   $(lsb_release -cs) \
   stable" && \
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - && \
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) && \
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list && \
sudo apt-get update && \
sudo apt-get -qq -y install nvidia-container-runtime nvidia-docker2 docker-ce && \
sudo curl -L https://github.com/docker/compose/releases/download/1.29.2/docker-compose-`uname -s`-`uname -m` -o /usr/local/bin/docker-compose && \
sudo chmod +x /usr/local/bin/docker-compose && \
sudo usermod -aG docker $USER && \
sudo docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi && \
sudo sysctl -w vm.max_map_count=262144 && \
exit

# You can validate if the GPUs are supported by running:
$ docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi