Update Dockerfile

Dockerfile

@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV MKL_THREADING_LAYER=GNU
@@ -6,17 +6,17 @@ ENV MKL_THREADING_LAYER=GNU
 # Install dependencies
 RUN apt-get update && apt-get install -y \
     git wget libgl1-mesa-glx libglib2.0-0 ffmpeg libx264-dev pkg-config \
-    build-essential cmake g++ python3-dev software-properties-common \
-    wget gnupg2 libopenblas-dev libx11-dev && \
+    build-essential cmake g++ software-properties-common \
+    gnupg2 libopenblas-dev libx11-dev && \
     rm -rf /var/lib/apt/lists/*
 
 # Add NVIDIA package repositories
-RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb && \
-    dpkg -i cuda-keyring_1.0-1_all.deb && \
+RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \
+    dpkg -i cuda-keyring_1.1-1_all.deb && \
     apt-get update
 
-# Install CUDA Toolkit and cuDNN
-RUN apt-get install -y cuda-toolkit-11-7 libcudnn8 libcudnn8-dev
+# Install CUDA Toolkit and cuDNN (with version matching the base image)
+RUN apt-get install -y cuda-toolkit-12-1 libcudnn8 libcudnn8-dev
 
 # Set CUDA paths for build environment
 ENV PATH=/usr/local/cuda/bin:$PATH \
@@ -24,7 +24,7 @@ ENV PATH=/usr/local/cuda/bin:$PATH \
     CUDA_HOME=/usr/local/cuda \
     LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
     CMAKE_PREFIX_PATH="/usr/local/cuda:/usr/lib/x86_64-linux-gnu/" \
-    TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6"
+    TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6;8.9;9.0"
 
 # Upgrade pip & install dependencies
 RUN pip install --upgrade pip setuptools wheel