Update Dockerfile

Dockerfile

@@ -13,18 +13,15 @@ RUN apt-get update && \
     gnupg \
     && rm -rf /var/lib/apt/lists/*
 
-# Add NVIDIA package repositories
-RUN wget https://developer.download.nvidia.com/compute/cuda/repos/debian10/x86_64/cuda-10-1_10.1.243-1_amd64.deb && \
-    dpkg -i cuda-10-1_10.1.243-1_amd64.deb && \
-    rm cuda-10-1_10.1.243-1_amd64.deb && \
+# Install CUDA Toolkit
+RUN wget https://developer.download.nvidia.com/compute/cuda/repos/debian10/x86_64/cuda-11-4_11.4.1-1_amd64.deb && \
+    dpkg -i cuda-11-4_11.4.1-1_amd64.deb && \
+    rm cuda-11-4_11.4.1-1_amd64.deb && \
     apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/debian10/x86_64/7fa2af80.pub && \
     echo "deb http://developer.download.nvidia.com/compute/cuda/repos/debian10/x86_64 /" > /etc/apt/sources.list.d/cuda.list && \
-    apt-get update
-
-# Install CUDA and other dependencies
-RUN apt-get install -y --no-install-recommends \
-    cuda-cudart-10-1 \
-    cuda-compat-10-1 \
+    apt-get update && \
+    apt-get install -y --no-install-recommends \
+    cuda-toolkit-11-4 \
     && rm -rf /var/lib/apt/lists/*
 
 # Install required Python packages