Update Dockerfile

Dockerfile

@@ -27,6 +27,20 @@ RUN chmod +x /start.sh
 RUN mkdir /data && chown user:user /data
 RUN ln -s /data /app
 
+# Install CUDA Toolkit 12.4 base installer and NVIDIA driver
+RUN apt-get update && apt-get install -y wget gnupg2 software-properties-common
+RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.1-1_all.deb
+RUN dpkg -i cuda-keyring_1.1-1_all.deb
+RUN apt-get update
+# Installing CUDA Toolkit 12.4
+RUN apt-get -y install cuda-toolkit-12-4
+# Installing NVIDIA Drivers (Legacy Kernel Module Flavor)
+RUN apt-get install -y cuda-drivers
+
+# Setting up environment variables for CUDA
+ENV PATH=/usr/local/cuda-12.4/bin:${PATH}
+ENV LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64:${LD_LIBRARY_PATH}
+
 # Switch back to the non-root user for running applications
 USER user
 ENTRYPOINT ["/start.sh"]