diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 7728975f95b50dc3a19d0519432f722e1840126b..3b20031d2eb209589c36b764dab3d0348e2ea071 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -27,23 +27,7 @@ before_script:
   - mkdir -p ~/.docker
   - echo $DOCKER_CONFIG_JSON | base64 -d > ~/.docker/config.json
   - docker system prune -a -f --volumes || true
-  - docker container prune -f || true
-  - docker network prune -f || true
-  - docker image prune -a -f || true
-  - docker builder prune -f || true
-  - apt-get update
-  - apt-get install -y apt-utils
-  - apt-get clean
-  - rm -rf /var/lib/apt/lists/*
-  - rm -rf /tmp/*
-  - journalctl --vacuum-time=1s
-
-after_script:
-  - docker system prune -a -f --volumes || true
-  - docker container prune -f || true
-  - docker network prune -f || true
-  - docker image prune -a -f || true
-  - docker builder prune -f || true
+  - apt-get update && apt-get install -y apt-utils && apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/*
 
 build:
   stage: build
@@ -52,7 +36,6 @@ build:
     - name: docker:20.10.16-dind
       command: ["--host=tcp://docker:2375", "--tls=false", "--storage-driver=overlay2", "--mtu=1500", "--data-root=/mnt/volume"]
   script:
-    - docker info
     - 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 push 22baar1mst/asyst-nlp-grader:latest
diff --git a/flask_ml_api/Dockerfile b/flask_ml_api/Dockerfile
index a09037025d8a4d34364584240da3b4fcafc880b9..8bfd6859383856549ceaf8b654c4f4bafe48cb30 100755
--- a/flask_ml_api/Dockerfile
+++ b/flask_ml_api/Dockerfile
@@ -1,37 +1,17 @@
-=FROM pytorch/pytorch:1.10.0-cuda11.3-cudnn8-runtime as build
-
-RUN apt-get update && apt-get install -y curl
-RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs -o /usr/local/bin/rustup-init.sh
-RUN chmod +x /usr/local/bin/rustup-init.sh
-RUN /usr/local/bin/rustup-init.sh -y --default-toolchain none
-
-ENV PATH="$PATH:/root/.cargo/bin"
-
-RUN /root/.cargo/bin/rustup toolchain install stable --profile minimal --no-self-update
-
-RUN /root/.cargo/bin/rustup --version
-RUN /root/.cargo/bin/cargo --version
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime as base
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    build-essential python3-venv && \
+    build-essential python3-venv curl && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
 RUN python3 -m venv /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
-
-FROM pytorch/pytorch:1.10.0-cuda11.3-cudnn8-runtime
+    pip install --no-cache-dir --upgrade setuptools wheel
 
-COPY --from=build /opt/myenv /opt/myenv
-
-ENV PATH="/opt/myenv/bin:$PATH"
+FROM base as build
 
 WORKDIR /app
 COPY . /app
@@ -39,5 +19,4 @@ COPY . /app
 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