Update Dockerfile

Dockerfile

@@ -1,11 +1,23 @@
 FROM python:3.9
 
+# Use an official TensorFlow GPU base image
+FROM tensorflow/tensorflow:latest-gpu
+
+# Set the working directory in the container
 WORKDIR /code
 
-COPY ./requirements.txt /code/requirements.txt
+# Install CUDA drivers (modify the version according to your requirements)
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    cuda-drivers \
+ && apt-get clean \
+ && rm -rf /var/lib/apt/lists/*
 
+# Copy requirements and install Python dependencies
+COPY ./requirements.txt /code/requirements.txt
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
+# Copy the rest of the application code
 COPY . .
 
-CMD ["python", "caption_api.py", "--host", "0.0.0.0", "--port", "7860"]
+# Set the default command to run your application
+CMD ["python", "caption_api.py", "--host", "0.0.0.0", "--port", "7860"]