1 files changed, 23 insertions, 3 deletions
diff --git a/src/nvidia-cuda/install.sh b/src/nvidia-cuda/install.sh
index f017ecb..9b41713 100644
--- a/src/nvidia-cuda/install.sh
+++ b/src/nvidia-cuda/install.sh
@@ -2,6 +2,9 @@
 
 set -e
 
+# Clean up
+rm -rf /var/lib/apt/lists/*
+
 INSTALL_CUDNN=${INSTALLCUDNN}
 INSTALL_NVTX=${INSTALLNVTX}
 CUDA_VERSION=${CUDAVERSION}
@@ -12,9 +15,23 @@ if [ "$(id -u)" -ne 0 ]; then
     exit 1
 fi
 
-# Install dependencies
-apt-get update -yq
-apt-get install -yq wget ca-certificates
+apt_get_update()
+{
+    if [ "$(find /var/lib/apt/lists/* | wc -l)" = "0" ]; then
+        echo "Running apt-get update..."
+        apt-get update -y
+    fi
+}
+
+# Checks if packages are installed and installs them if not
+check_packages() {
+    if ! dpkg -s "$@" > /dev/null 2>&1; then
+        apt_get_update
+        apt-get -y install --no-install-recommends "$@"
+    fi
+}
+
+check_packages wget ca-certificates
 
 # Add NVIDIA's package repository to apt so that we can download packages
 # Always use the ubuntu2004 repo because the other repos (e.g., debian11) are missing packages
@@ -55,4 +72,7 @@ if [ "$INSTALL_NVTX" = "true" ]; then
     apt-get install -yq "$nvtx_pkg"
 fi
 
+# Clean up
+rm -rf /var/lib/apt/lists/*
+
 echo "Done!"