diff --git a/vm-images/elements/cuda/post-install.d/05-cuda-install b/vm-images/elements/cuda/post-install.d/05-cuda-install
index 503c2f6..0b3af57 100755
--- a/vm-images/elements/cuda/post-install.d/05-cuda-install
+++ b/vm-images/elements/cuda/post-install.d/05-cuda-install
@@ -31,9 +31,12 @@ echo 'debconf debconf/frontend select Noninteractive' | sudo debconf-set-selecti
 DEBIAN_FRONTEND=noninteractive sudo apt-get install -y -q
 DEBIAN_FRONTEND=noninteractive sudo apt update
 DEBIAN_FRONTEND=noninteractive sudo apt upgrade -y
-DEBIAN_FRONTEND=noninteractive sudo apt install nvidia-driver-525 -y
+DEBIAN_FRONTEND=noninteractive sudo apt install -y ubuntu-drivers-common
+# The --gpgpu flag is required for server drivers (-server suffix)
+# See: https://documentation.ubuntu.com/server/how-to/graphics/install-nvidia-drivers/
+DEBIAN_FRONTEND=noninteractive sudo ubuntu-drivers install --gpgpu nvidia:580-server
 echo "Put nvidia drivers on hold: DO NOT UPGRADE basically"
-dpkg-query -W --showformat='${Package} ${Status}\n' | grep -v deinstall | awk '{ print $1 }' | grep -E 'nvidia.*-[0-9]+$' | xargs -r -L 1 sudo apt-mark hold
+dpkg-query -W --showformat='${Package} ${Status}\n' | grep -v deinstall | awk '{ print $1 }' | grep -E '(nvidia|libnvidia|linux-.+-nvidia)-.*580' | xargs -r -L 1 sudo apt-mark hold
 
 echo "#############################################################"
 echo "# Install Nvidia Container Toolkit                          #"