diff --git a/gpu_benchmark/gpu_setup.sh b/gpu_benchmark/gpu_setup.sh
index 8b3289a8..0798bec5 100755
--- a/gpu_benchmark/gpu_setup.sh
+++ b/gpu_benchmark/gpu_setup.sh
@@ -1,6 +1,7 @@
 #!/bin/bash
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2023 United Kingdom Research and Innovation
+set -ex
 
 sudo dnf update -y
 
@@ -13,6 +14,9 @@ sudo dracut --force
 
 sudo dnf install tar bzip2 make automake gcc gcc-c++ pciutils elfutils-libelf-devel libglvnd-devel -y
 sudo dnf install -y kernel-devel kernel-headers -y
-wget -nc https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run
-nvidia-smi || sudo sh cuda_12.1.0_530.30.02_linux.run --silent
-nvidia-smi || (echo "Rebooting machine to load Nvidia Driver" && sudo reboot)
\ No newline at end of file
+
+VERSION=$(awk -F= '/^VERSION_ID=/ {gsub("\"","",$2); print $2}' /etc/os-release 2>/dev/null || true)
+sudo dnf config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel"${VERSION%%.*}"/"$(uname -m)"/cuda-rhel"${VERSION%%.*}".repo
+
+nvidia-smi || sudo dnf module install nvidia-driver:580 -y
+nvidia-smi || (echo "Rebooting machine to load Nvidia Driver" && sudo reboot)