diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 3b20031d2eb209589c36b764dab3d0348e2ea071..b1551ba73c34a99c4119786186efd30ed61fac59 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -23,30 +23,38 @@ variables:
   DOCKER_CPUS: "2"
   DOCKER_REGISTRY: docker.io
 
-before_script:
-  - mkdir -p ~/.docker
-  - echo $DOCKER_CONFIG_JSON | base64 -d > ~/.docker/config.json
-  - docker system prune -a -f --volumes || true
-  - apt-get update && apt-get install -y apt-utils && apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/*
-
 build:
   stage: build
-  image: docker:20.10.16
+  image: docker:dind
   services:
-    - name: docker:20.10.16-dind
-      command: ["--host=tcp://docker:2375", "--tls=false", "--storage-driver=overlay2", "--mtu=1500", "--data-root=/mnt/volume"]
+    - name: docker:dind
+      command: ["--tls=false", "--storage-driver=overlay2", "--mtu=1500", "--data-root=/mnt/volume"]
+  before_script:
+    - mkdir -p ~/.docker
+    - echo $DOCKER_CONFIG_JSON | base64 -d > ~/.docker/config.json
+    - mkdir -p /mnt/volume
+    - mount /dev/vdb /mnt/volume || echo "Failed to mount /dev/vdb"
+    - export DOCKER_TMPDIR=/mnt/volume
+    - docker system prune -a -f --volumes || true
+    - df -h
+  after_script:
+    - docker system prune -a -f
+    - docker volume prune -f
   script:
-    - docker buildx create --use
-    - 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 .
+    - docker build --cache-from=asyst-nlp-grader --no-cache --progress=plain -t asyst-nlp-grader -f flask_ml_api/Dockerfile .
     - docker push 22baar1mst/asyst-nlp-grader:latest
-  only:
-    - master
+  cache:
+    key: build-cache
+    paths:
+      - /mnt/volume
+  timeout: 1h
 
 deploy:
   stage: deploy
-  image: docker:20.10.16
+  dependencies:
+    - build
   services:
-    - name: docker:20.10.16-dind
+    - name: docker:dind
       command: ["--tls=false", "--storage-driver=overlay2", "--mtu=1500", "--data-root=/mnt/volume"]
   script:
     - docker pull 22baar1mst/asyst-nlp-grader:latest
diff --git a/flask_ml_api/Dockerfile b/flask_ml_api/Dockerfile
index ab1292818b2d885883f8ce38860b0a395f4741b9..c1ca87072026ccf4ab613974fb9129538908be26 100755
--- a/flask_ml_api/Dockerfile
+++ b/flask_ml_api/Dockerfile
@@ -1,31 +1,30 @@
-FROM debian:bullseye-slim as base
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime AS builder
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    build-essential python3-venv curl && \
+    build-essential python3-venv && \
     apt-get clean && \
-    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* /usr/share/doc /usr/share/man /usr/share/locale
+    rm -rf /var/lib/apt/lists/*
 
 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 \
-        rm -rf /root/.cache/pip
+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
 
-FROM base as build
-COPY --from=build /opt/myenv /opt/myenv
+# Final image
+FROM python:3.10-slim
 
-WORKDIR /app
-COPY . /app
+COPY --from=builder /opt/myenv /opt/myenv
+COPY --from=builder /app /app
+
+ENV PATH="/opt/myenv/bin:$PATH"
 
 RUN chown -R www-data:www-data /app && chmod -R 755 /app
 
 EXPOSE 5000
-CMD ["/opt/myenv/bin/python", "/app/api.py"]
\ No newline at end of file
+
+CMD ["/opt/myenv/bin/python", "/app/api.py"]