Update Dockerfile

Dockerfile

@@ -10,8 +10,21 @@ COPY . .
 # Install any needed packages specified in requirements.txt
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
-    cuda-cudart-11-4=11.4.1-1 \
-    cuda-compat-11-4 \
+    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 && \
+    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 \
     && rm -rf /var/lib/apt/lists/*
 
 # Install required Python packages