cuda_setup.sh

#!/usr/bin/env bash
set -euxo pipefail

export NVIDIA_DRIVER_VERSION=470-server
export CUDA_VERSION=11-1

modify_bashrc() {
  # don't variable expand in heredoc
  cat << 'EOF' >> /etc/profile.d/cuda.sh
export PATH=$PATH:/usr/local/cuda/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64
EOF
}

setup_repos() {
  # shellcheck disable=SC1091
  source /etc/os-release
  ARCH="$(arch)"
  VERSION_ID_FIXED=${VERSION_ID//./}

  # add CUDA and cuDNN repository
  # for WSL: use https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/
  curl -fsSL "https://developer.download.nvidia.com/compute/cuda/repos/$ID$VERSION_ID_FIXED/$ARCH/3bf863cc.pub" |
    sudo gpg --no-default-keyring --keyring /usr/share/keyrings/nvidia-cuda.gpg --import
  curl -fsSL "https://developer.download.nvidia.com/compute/machine-learning/repos/$ID$VERSION_ID_FIXED/$ARCH/7fa2af80.pub" |
    sudo gpg --no-default-keyring --keyring /usr/share/keyrings/nvidia-ml.gpg --import
  echo "deb [signed-by=/usr/share/keyrings/nvidia-cuda.gpg] https://developer.download.nvidia.com/compute/cuda/repos/$ID$VERSION_ID_FIXED/$ARCH /" |
    sudo tee /etc/apt/sources.list.d/cuda.list
  echo "deb [signed-by=/usr/share/keyrings/nvidia-ml.gpg] https://developer.download.nvidia.com/compute/machine-learning/repos/$ID$VERSION_ID_FIXED/$ARCH /" |
    sudo tee /etc/apt/sources.list.d/nvidia-ml.list

  # add nvidia-container-toolkit repository
  curl -fsSL https://nvidia.github.io/nvidia-docker/gpgkey |
    sudo gpg --no-default-keyring --keyring /usr/share/keyrings/nvidia-container.gpg --import
  curl -fsSL "https://nvidia.github.io/nvidia-docker/$ID$VERSION_ID/nvidia-docker.list" |
    sed "s/deb https/deb [signed-by=\/usr\/share\/keyrings\/nvidia-container.gpg] https/g" |
    sudo tee /etc/apt/sources.list.d/nvidia-container.list
  
  sudo apt update
  sudo apt upgrade --yes
}

setup_driver() {
  setup_repos
  
  # install requirements
  sudo apt install --yes \
    "linux-headers-$(uname -r)" \
    build-essential

  # NVIDIA driver only
  sudo apt install --yes --no-install-recommends \
    nvidia-driver-"$NVIDIA_DRIVER_VERSION" \
    nvidia-container-toolkit
    
  echo 'NVIDIA driver setup done.'
  echo 'Reboot now.'
}

setup_cuda() {
  setup_repos

  # no GPU driver
  sudo apt install --yes --no-install-recommends \
    cuda-toolkit-"$CUDA_VERSION"

  modify_bashrc
  
  echo 'CUDA setup done.'
}

echo 'Setup CUDA...'
eval "$1"

NOTE: This script currently supports only for Ubuntu distro.