diff --git a/docker/scripts/install_flash_attension.sh b/docker/scripts/install_flash_attension.sh
index 6a3301c2..f37e567d 100644
--- a/docker/scripts/install_flash_attension.sh
+++ b/docker/scripts/install_flash_attension.sh
@@ -1,6 +1,4 @@
-    git clone -b v1.0.8 https://github.com/Dao-AILab/flash-attention && \
-    cd flash-attention && pip install . && \
-    pip install csrc/layer_norm && \
-    pip install csrc/rotary && \
+    git clone -b v2.3.2 https://github.com/Dao-AILab/flash-attention && \
+    cd flash-attention && python setup.py install && \
     cd .. && \
     rm -rf flash-attention