diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index a146540ccf2d9ec686d80f8abf9dfbb20168db5d..5da8f94e059a5d3d04f7b1972f5f82d1c2b01424 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -42,8 +42,8 @@ build:
   script:
     - docker buildx create --use
     #    - docker build --cache-from=flask-backend-full --no-cache --progress=plain -t flask-backend-full -f flask_ml_api/Dockerfile .
-    - docker build --cache-from=flask-backend-full --no-cache --progress=plain -t 22baar1mst/asyst-nlp-grader:latest -f flask_ml_api/Dockerfile .
-    - docker build --no-cache --progress=plain -t 22baar1mst/asyst-nlp-grader:latest -f flask_ml_api/Dockerfile .
+    - docker build --cache-from=22baar1mst/asyst-nlp-grader --no-cache --progress=plain -t 22baar1mst/asyst-nlp-grader:latest -f flask_ml_api/Dockerfile .
+#    - docker build --no-cache --progress=plain -t 22baar1mst/asyst-nlp-grader:latest -f flask_ml_api/Dockerfile .
     - echo $DOCKER_REPO
     - docker push 22baar1mst/asyst-nlp-grader:latest
   only:
diff --git a/flask_ml_api/Dockerfile b/flask_ml_api/Dockerfile
index 4c16d3f9afb899b0c151c6b85f78eefd13e45543..49dd2be619a905bbaaa8c04883e9e478095245ba 100755
--- a/flask_ml_api/Dockerfile
+++ b/flask_ml_api/Dockerfile
@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime as build
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     build-essential python3-venv && \
@@ -8,13 +8,26 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 RUN python3 -m venv /opt/myenv
 ENV PATH="/opt/myenv/bin:$PATH"
 
+#RUN pip install --no-cache-dir Flask matplotlib scikit-learn transformers pandas sentence_transformers && \
+#    pip install --no-cache-dir --upgrade setuptools wheel && \
+#    rm -rf /root/.cache/pip
+
+RUN pip install --no-cache-dir Flask && \
+    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
+
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+COPY --from=build /opt/myenv /opt/myenv
+
+ENV PATH="/opt/myenv/bin:$PATH"
+
 WORKDIR /app
 COPY . /app
 
-RUN pip install --no-cache-dir Flask matplotlib scikit-learn transformers pandas sentence_transformers && \
-    pip install --no-cache-dir --upgrade setuptools wheel && \
-    rm -rf /root/.cache/pip
-
 RUN chown -R www-data:www-data /app && chmod -R 755 /app
 
 EXPOSE 5000