diff --git a/docker/Dockerfile b/docker/Dockerfile
index 201187c..4c6033f 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -19,9 +19,10 @@ WORKDIR /workspace/
 COPY ./requirements.txt .
 RUN pip install --upgrade pip \
     && pip install -r requirements.txt \
+    && pip install flash-attn==2.6.0.post1 --no-build-isolation \
     && pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary \
     && pip install causal-conv1d==1.4.0 \
-    && pip install mamba-ssm==2.2.2
+    && pip install mamba-ssm==2.2.2 
     
 RUN [ "python3", "-c", "import nltk; nltk.download('punkt')"]
 
diff --git a/docker/requirements.txt b/docker/requirements.txt
index 6c47aa9..f518f81 100644
--- a/docker/requirements.txt
+++ b/docker/requirements.txt
@@ -12,6 +12,5 @@ tiktoken
 tenacity
 accelerate
 huggingface_hub==0.23.4
-transformers==4.44.1
-flash-attn==2.6.0.post1
+transformers==4.44.2
 vllm==0.5.4
\ No newline at end of file