Update Dockerfile

Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
+FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
@@ -7,7 +7,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
     LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
     TORCH_CUDA_ARCH_LIST="8.6" \
     NVIDIA_VISIBLE_DEVICES=all \
-    NVIDIA_DRIVER_CAPABILITIES=compute,utility,graphics
+    NVIDIA_DRIVER_CAPABILITIES=all
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     build-essential \
@@ -17,10 +17,21 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     ffmpeg \
     libsndfile1 \
     curl \
-    nvidia-cuda-toolkit \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
+# Python environment configuration
+RUN ln -sf /usr/bin/python3 /usr/bin/python && \
+    curl -sS https://bootstrap.pypa.io/get-pip.py | python3 && \
+    python3 -m pip install --no-cache-dir pip==23.3.1 setuptools==69.0.3 wheel==0.42.0
+
+# Install PyTorch with CUDA support
+RUN pip3 install --no-cache-dir torch==2.1.2 torchvision==0.16.2 --index-url https://download.pytorch.org/whl/cu121
+
+WORKDIR /app
+RUN mkdir -p /app/cache /app/src && \
+    chmod 777 /app/cache
+
 # Rest of your Dockerfile remains the same...
 # # System-level configuration
 # ENV DEBIAN_FRONTEND=noninteractive \