Add CUDA 12.8 libtorch image

tinglvv · tinglvv · commit 8196cbd4d69b · 2025-01-27T14:50:26.000-08:00
diff --git a/.ci/docker/libtorch/Dockerfile b/.ci/docker/libtorch/Dockerfile
@@ -56,11 +56,6 @@ RUN bash ./install_cuda.sh 11.8
 RUN bash ./install_magma.sh 11.8
 RUN ln -sf /usr/local/cuda-11.8 /usr/local/cuda
 
-FROM cuda as cuda12.1
-RUN bash ./install_cuda.sh 12.1
-RUN bash ./install_magma.sh 12.1
-RUN ln -sf /usr/local/cuda-12.1 /usr/local/cuda
-
 FROM cuda as cuda12.4
 RUN bash ./install_cuda.sh 12.4
 RUN bash ./install_magma.sh 12.4
@@ -71,6 +66,11 @@ RUN bash ./install_cuda.sh 12.6
 RUN bash ./install_magma.sh 12.6
 RUN ln -sf /usr/local/cuda-12.6 /usr/local/cuda
 
+FROM cuda as cuda12.8
+RUN bash ./install_cuda.sh 12.8
+RUN bash ./install_magma.sh 12.8
+RUN ln -sf /usr/local/cuda-12.8 /usr/local/cuda
+
 FROM cpu as rocm
 ARG PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
diff --git a/.github/workflows/build-libtorch-images.yml b/.github/workflows/build-libtorch-images.yml
@@ -45,7 +45,7 @@ jobs:
     runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.9xlarge.ephemeral"
     strategy:
       matrix:
-        cuda_version: ["12.6", "12.4", "11.8"]
+        cuda_version: ["12.8", "12.6", "12.4", "11.8"]
     env:
       GPU_ARCH_TYPE: cuda
       GPU_ARCH_VERSION: ${{ matrix.cuda_version }}
