From d4f6a6b1032b87061c760f558093a9168dfe77c7 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Tue, 30 Jul 2024 08:34:37 -0400
Subject: [PATCH] fix dockerfile and base builder (#1795) [skip-ci]

---
 .github/workflows/base.yml | 4 ++++
 docker/Dockerfile-base     | 2 +-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml
index 4019a5baf6..3a0c143df5 100644
--- a/.github/workflows/base.yml
+++ b/.github/workflows/base.yml
@@ -14,16 +14,19 @@ jobs:
         include:
           - cuda: "121"
             cuda_version: 12.1.1
+            cudnn_version: 8
             python_version: "3.10"
             pytorch: 2.3.1
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
           - cuda: "121"
             cuda_version: 12.1.1
+            cudnn_version: 8
             python_version: "3.11"
             pytorch: 2.3.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.0
+            cudnn_version: ""
             python_version: "3.11"
             pytorch: 2.4.0
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
@@ -52,6 +55,7 @@ jobs:
           labels: ${{ steps.metadata.outputs.labels }}
           build-args: |
             CUDA_VERSION=${{ matrix.cuda_version }}
+            CUDNN_VERSION=${{ matrix.cudnn_version }}
             CUDA=${{ matrix.cuda }}
             PYTHON_VERSION=${{ matrix.python_version }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
index 1de5537dac..3f13bba30a 100644
--- a/docker/Dockerfile-base
+++ b/docker/Dockerfile-base
@@ -3,7 +3,7 @@ ARG CUDNN_VERSION="8"
 ARG UBUNTU_VERSION="22.04"
 ARG MAX_JOBS=4
 
-FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION as base-builder
+FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder
 
 ENV PATH="/root/miniconda3/bin:${PATH}"