diff --git a/ci/build_conda.sh b/ci/build_conda.sh
index a4bc9816..d0fbfbf9 100755
--- a/ci/build_conda.sh
+++ b/ci/build_conda.sh
@@ -10,10 +10,7 @@ export CMAKE_GENERATOR=Ninja
 rapids-print-env
 
 rapids-logger "Install CUDA Toolkit"
-yum update
-yum install -y epel-release
-yum install -y nvidia-driver-latest-dkms
-yum install -y cuda-toolkit-12-3
+source "$(dirname "$0")/install_latest_cuda_toolkit.sh"
 
 rapids-logger "Begin py build"
 
diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index 4604e5ef..41237b61 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -3,10 +3,7 @@
 
 
 rapids-logger "Install CUDA Toolkit"
-yum update
-yum install -y epel-release
-yum install -y nvidia-driver-latest-dkms
-yum install -y cuda-toolkit-12-3
+source "$(dirname "$0")/install_latest_cuda_toolkit.sh"
 
 rapids-logger "Check conda environment"
 conda list
diff --git a/ci/install_latest_cuda_toolkit.sh b/ci/install_latest_cuda_toolkit.sh
new file mode 100755
index 00000000..d5232de5
--- /dev/null
+++ b/ci/install_latest_cuda_toolkit.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION
+
+# Installs the latest CUDA Toolkit.
+# Supports CentOS 7 and Rocky Linux 8.
+
+yum update
+yum install -y epel-release
+
+OS_ID=$(. /etc/os-release; echo $ID)
+if [ "${OS_ID}" == "centos" ]; then
+    yum install -y nvidia-driver-latest-dkms
+elif [ "${OS_ID}" == "rocky" ]; then
+    yum install -y nvidia-driver
+else
+    echo "Error: OS not detected as CentOS or Rocky Linux. Exiting."
+    exit 1
+fi
+
+yum install -y cuda-toolkit-12-3