From 9ee7ce5c85a71b4d9c2b6f0fab382ca21708e88e Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 29 Oct 2025 12:59:26 -0400 Subject: [PATCH] set TORCH_CUDA_ARCH_LIST correctly --- cicd/Dockerfile.jinja | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cicd/Dockerfile.jinja b/cicd/Dockerfile.jinja index 6a1ddb66d..c3a613ecc 100644 --- a/cicd/Dockerfile.jinja +++ b/cicd/Dockerfile.jinja @@ -1,6 +1,6 @@ FROM axolotlai/axolotl-base:{{ BASE_TAG }} -ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX" +ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX" ENV AXOLOTL_EXTRAS="{{ AXOLOTL_EXTRAS }}" ENV AXOLOTL_ARGS="{{ AXOLOTL_ARGS }}" ENV CUDA="{{ CUDA }}"