diff --git a/Dockerfile b/Dockerfile
index e0c805c..03b7057 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,34 +1,22 @@
-# Use the official Python 3.10 image as base
-FROM python:3.10-slim
+FROM pytorch/pytorch:2.2.1-cuda12.1-cudnn8-devel
 
 # Set environment variables
-ENV PYTHONUNBUFFERED=1
 WORKDIR /app
 
-# Install required system packages
-RUN apt-get update && apt-get install -y \
+# Install required system packages and Python
+RUN apt-get update -y && apt-get install -y \
     curl \
-    wget \
-    git \
-    sudo \
-    && rm -rf /var/lib/apt/lists/*
+    git
+# Custom Lean 4 installation script without Visual Studio Code
+RUN curl -y https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -sSf | sh
 
-# Install Lean 4
-RUN wget -q https://raw.githubusercontent.com/leanprover-community/mathlib4/master/scripts/install_debian.sh \
-    && bash install_debian.sh \
-    && rm -f install_debian.sh \
-    && source ~/.profile
-
-# Add lean to the system's PATH
-ENV PATH="/root/.elan/bin:$PATH"
-
-# Install Python dependencies (assumes that the requirements.txt is in the repository)
+# Copy the requirements.txt file
 COPY requirements.txt /app/
-RUN pip install --no-cache-dir -r requirements.txt
+RUN pip3 install --no-cache-dir -r requirements.txt
 
 # Build Mathlib4 during the image build
-WORKDIR /app/mathlib4
-RUN lake build
+# WORKDIR /app/mathlib4
+# RUN lake exe cache get && lake build
 
 # Set default work directory back to project root
 WORKDIR /app
@@ -37,3 +25,4 @@ WORKDIR /app
 # EXPOSE 8080
 
 CMD ["bash"]
+
diff --git a/requirements.txt b/requirements.txt
index 53f7e93..cd600ab 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
+torch==2.2.1
 pytz==2022.1
 easydict==1.13
-torch==2.2.1
 transformers==4.40.1
 vllm==0.4.1
 numpy==1.26.4
@@ -8,4 +8,4 @@ pandas==1.4.3
 tabulate==0.9.0
 termcolor==2.4.0
 accelerate==0.33.0
-flash_attn==2.6.3
\ No newline at end of file
+flash_attn==2.6.3