Image build is optimized.

flask_ml_api/Dockerfile

-# Build stage
-FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime AS builder
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
 
+# Install necessary packages
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    build-essential python3-venv && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+build-essential python3-venv && \
+apt-get clean && \
+rm -rf /var/lib/apt/lists/*
+
+# Create virtual environment and install dependencies
+RUN python3 -m venv /opt/myenv && \
+/opt/myenv/bin/pip install --no-cache-dir Flask matplotlib scikit-learn transformers pandas sentence_transformers && \
+/opt/myenv/bin/pip install --no-cache-dir --upgrade setuptools wheel
 
-RUN python3 -m venv /opt/myenv
 ENV PATH="/opt/myenv/bin:$PATH"
 
 WORKDIR /app
 COPY . /app
 
-# Install dependencies in smaller steps
-RUN pip install --no-cache-dir Flask matplotlib scikit-learn transformers pandas sentence_transformers
-RUN pip install --no-cache-dir --upgrade setuptools wheel
-
-# Final image
-FROM python:3.10-slim
-
-COPY --from=builder /opt/myenv /opt/myenv
-COPY --from=builder /app /app
-
-ENV PATH="/opt/myenv/bin:$PATH"
-
+# Setting permissions and cleaning
 RUN chown -R www-data:www-data /app && chmod -R 755 /app
 
 EXPOSE 5000
 
-CMD ["/opt/myenv/bin/python", "/app/api.py"]
+CMD ["/opt/myenv/bin/python", "/app/api.py"]
\ No newline at end of file