lingyuzeng/cdc_dockerfile
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update TORCH_CUDA_ARCH_LIST

Browse Source
This commit is contained in:
lingyuzeng
2024-07-15 12:25:54 +08:00
parent 3fc430582e
commit 47f460bbc0
2 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
finetune/Dockerfile.update
View File

@@ -110,8 +110,8 @@ python3 -m pip uninstall -y torch-tensorrt
python3 -m pip uninstall -y apex
EOT
# install apex
ARG TORCH_CUDA_ARCH_LIST="6.0;6.1;6.2;7.0;7.5;8.0;8.6;8.9,9.0"
# install apex TORCH_CUDA_ARCH_LIST all "6.0;6.1;6.2;7.0;7.5;8.0;8.6;8.9;9.0"
ARG TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0"
ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
RUN <<EOT
#!/bin/bash