From 1125bb11cbdd191ce8e116f68015c5733931580a Mon Sep 17 00:00:00 2001
From: udaij12 <32673964+udaij12@users.noreply.github.com>
Date: Tue, 9 Jul 2024 13:14:51 -0700
Subject: [PATCH] Updating docker cuda and github branch (#3233)

* cuda update and branch fix

* adding print

* testing cuda

* testing branch

---------

Co-authored-by: Ankith Gunapal <agunapal@ischool.Berkeley.edu>
---
 .github/workflows/docker-ci.yaml | 12 +++++++++++-
 docker/build_image.sh            |  4 ++--
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/docker-ci.yaml b/.github/workflows/docker-ci.yaml
index 6ad10a8713..0e7db91825 100644
--- a/.github/workflows/docker-ci.yaml
+++ b/.github/workflows/docker-ci.yaml
@@ -24,12 +24,22 @@ jobs:
         working-directory: docker
         run: ./test_build_image_tagging.sh ${{ matrix.python-version }}
 
+      - name: Determine the branch name
+        id: branch-name
+        run: |
+          if [[ "${GITHUB_REF}" == refs/heads/* ]]; then
+            GITHUB_BRANCH=${GITHUB_REF#refs/heads/}
+          elif [[ "${GITHUB_REF}" == refs/pull/*/merge ]]; then
+            GITHUB_BRANCH=${GITHUB_HEAD_REF}
+          fi
+          echo "GITHUB_BRANCH=${GITHUB_BRANCH}" >> $GITHUB_OUTPUT
+
       - name: Build Image for container test
         id: image_build
         working-directory: docker
         run: |
           IMAGE_TAG=test-image-${{ matrix.python-version }}
-          ./build_image.sh -py "${{ matrix.python-version }}" -t "${IMAGE_TAG}" -b ${GITHUB_HEAD_REF} -s
+          ./build_image.sh -py "${{ matrix.python-version }}" -t "${IMAGE_TAG}" -b ${{ steps.branch-name.outputs.GITHUB_BRANCH }} -s
           echo "IMAGE_TAG=${IMAGE_TAG}" >> $GITHUB_OUTPUT
 
       - name: Container Healthcheck
diff --git a/docker/build_image.sh b/docker/build_image.sh
index 6ca105ec9f..4738021a8a 100755
--- a/docker/build_image.sh
+++ b/docker/build_image.sh
@@ -54,8 +54,8 @@ do
         -g|--gpu)
           MACHINE=gpu
           DOCKER_TAG="pytorch/torchserve:latest-gpu"
-          BASE_IMAGE="nvidia/cuda:11.8.0-base-ubuntu20.04"
-          CUDA_VERSION="cu117"
+          BASE_IMAGE="nvidia/cuda:12.1.1-base-ubuntu20.04"
+          CUDA_VERSION="cu121"
           shift
           ;;
         -bi|--baseimage)