Update Dockerfile

Dockerfile

@@ -1,4 +1,4 @@
-# Change this line to use a CUDA 11.8 base image
+# Use the CUDA 11.8 base image
 FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
 
 # Set up the environment
@@ -19,7 +19,11 @@ ENV HF_HOME=/app/cache
 # Create the cache directory and give it open permissions
 RUN mkdir -p /app/cache && chmod 777 /app/cache
 
-# Copy and install Python dependencies
+# ⭐⭐⭐ ADD THIS BLOCK ⭐⭐⭐
+# Install PyTorch and related libraries for CUDA 11.8 separately
+RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
+
+# Copy and install the rest of the Python dependencies
 COPY requirements.txt .
 RUN pip3 install --no-cache-dir -r requirements.txt