update flash-attn to 2.6.1 (#62)

update flash-atto
dptech-corp · Jul 17, 2024 · c62cfe8 · c62cfe8
1 parent ba5962e
commit c62cfe8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/docker/rdma/Dockerfile b/docker/rdma/Dockerfile
@@ -66,7 +66,7 @@ RUN cd /tmp && \
 
 RUN pip3 install --no-cache-dir tokenizers lmdb biopython ml-collections timeout-decorator urllib3 tree dm-tree && rm -rf ~/.cache/pip 
 
-RUN MAX_JOBS=1 pip3 install -U 'flash-attn<2.5.0' --no-build-isolation --no-cache-dir
+RUN MAX_JOBS=1 pip3 install -U 'flash-attn==2.6.1' --no-build-isolation --no-cache-dir
 
 RUN ldconfig && \
     apt-get clean && \