py310, fix cuda arg in deepspeed

2023-05-30 18:02:34 -04:00
parent 13b80937f9
commit c43c5c84ff
2 changed files with 11 additions and 1 deletions
--- a/docker/Dockerfile-base
+++ b/docker/Dockerfile-base
@@ -52,6 +52,8 @@ RUN git clone https://github.com/HazyResearch/flash-attention.git && \

 FROM base-builder AS deepspeed-builder

+ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
+
 WORKDIR /workspace

 RUN git clone https://github.com/microsoft/DeepSpeed.git && \