Image build is optimized.

flask_ml_api/Dockerfile

-FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime as base
+FROM debian:bullseye-slim as base
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     build-essential python3-venv curl && \
     apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* /usr/share/doc /usr/share/man /usr/share/locale
 
 RUN python3 -m venv /opt/myenv
 ENV PATH="/opt/myenv/bin:$PATH"
 
 RUN pip install --no-cache-dir Flask && \
+    pip install --no-cache-dir pytorch && \
     pip install --no-cache-dir matplotlib && \
     pip install --no-cache-dir scikit-learn && \
     pip install --no-cache-dir transformers && \
     pip install --no-cache-dir pandas && \
     pip install --no-cache-dir sentence_transformers && \
-    pip install --no-cache-dir --upgrade setuptools wheel
+    pip install --no-cache-dir --upgrade setuptools wheel \
+        rm -rf /root/.cache/pip
+
 
 FROM base as build
+COPY --from=build /opt/myenv /opt/myenv
 
 WORKDIR /app
 COPY . /app