Update Dockerfile

Dockerfile

@@ -1,8 +1,7 @@
-FROM python:3.9
+# Use an official CUDA image as the base image
+FROM nvidia/cuda:11.4.1-cudnn8-runtime-ubuntu20.04
 
-FROM nvidia/cuda:11.4.1-cudnn8-devel-ubuntu20.04
-
-# Install Python
+# Install Python and pip
 RUN apt-get update && \
     apt-get install -y python3.9 python3-pip
 
@@ -18,9 +17,8 @@ COPY ./requirements.txt /code/requirements.txt
 # Install dependencies
 RUN pip3 install --no-cache-dir --upgrade -r /code/requirements.txt
 
-
-# Copy the rest of the application code
+# Copy the application files to the container
 COPY . .
 
-# Set the default command to run your application
+# Specify the command to run on container start
 CMD ["python", "caption_api.py", "--host", "0.0.0.0", "--port", "7860"]