From a971eb4ce6166ccb2b73f6a89bb4462f65003178 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Wed, 5 Feb 2025 09:24:02 -0500
Subject: [PATCH] Torch 2.6 support for base docker image (#2312)

---
 .github/workflows/base.yml | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml
index 0f4b3c9ce..41b3a7c84 100644
--- a/.github/workflows/base.yml
+++ b/.github/workflows/base.yml
@@ -40,6 +40,12 @@ jobs:
             python_version: "3.11"
             pytorch: 2.5.1
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+          - cuda: "124"
+            cuda_version: 12.4.1
+            cudnn_version: ""
+            python_version: "3.11"
+            pytorch: 2.6.0
+            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
     steps:
       - name: Checkout
         uses: actions/checkout@v4