diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base index be67e8eb4..6862569b9 100644 --- a/docker/Dockerfile-base +++ b/docker/Dockerfile-base @@ -40,7 +40,7 @@ ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX" RUN git clone https://github.com/Dao-AILab/flash-attention.git && \ cd flash-attention && \ - git checkout v2.0.0 && \ + git checkout 9ee0ff1 && \ python3 setup.py bdist_wheel && \ cd csrc/fused_dense_lib && \ python3 setup.py bdist_wheel && \