diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index b1551ba73c34a99c4119786186efd30ed61fac59..a40fa5f213669056bdcfb6b619f1a92c26283af6 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -41,7 +41,8 @@ build:
     - docker system prune -a -f
     - docker volume prune -f
   script:
-    - docker build --cache-from=asyst-nlp-grader --no-cache --progress=plain -t asyst-nlp-grader -f flask_ml_api/Dockerfile .
+#    - docker build --cache-from=asyst-nlp-grader --no-cache --progress=plain -t asyst-nlp-grader -f flask_ml_api/Dockerfile .
+    - docker buildx build --cache-from=type=registry,ref=22baar1mst/asyst-nlp-grader --no-cache --progress=plain -t 22baar1mst/asyst-nlp-grader:latest -f flask_ml_api/Dockerfile --push .
     - docker push 22baar1mst/asyst-nlp-grader:latest
   cache:
     key: build-cache
diff --git a/flask_ml_api/Dockerfile b/flask_ml_api/Dockerfile
index c1ca87072026ccf4ab613974fb9129538908be26..d6fa7065c591e1e60f3fedea445d1eae8ea90d8e 100755
--- a/flask_ml_api/Dockerfile
+++ b/flask_ml_api/Dockerfile
@@ -11,12 +11,10 @@ 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
+FROM python:3.10-slim AS final
 
 COPY --from=builder /opt/myenv /opt/myenv
 COPY --from=builder /app /app