Update Dockerfile

Dockerfile

@@ -1,20 +1,23 @@
 FROM python:3.9
 
-# Use an official TensorFlow GPU base image
-FROM tensorflow/tensorflow:latest-gpu
+FROM nvidia/cuda:11.4.1-cudnn8-devel-ubuntu20.04
+
+# Install Python
+RUN apt-get update && \
+    apt-get install -y python3.9 python3-pip
+
+# Install NumPy and GRPCIO
+RUN pip3 install --no-cache-dir --upgrade numpy grpcio
 
 # Set the working directory in the container
 WORKDIR /code
 
-# 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 the requirements file
 COPY ./requirements.txt /code/requirements.txt
-RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# Install dependencies
+RUN pip3 install --no-cache-dir --upgrade -r /code/requirements.txt
+
 
 # Copy the rest of the application code
 COPY . .