diff --git a/.github/workflows/build_and_push_ci_image.yml b/.github/workflows/build_and_push_ci_image.yml
index 4bb54ce0..12bb20ca 100644
--- a/.github/workflows/build_and_push_ci_image.yml
+++ b/.github/workflows/build_and_push_ci_image.yml
@@ -32,11 +32,17 @@ jobs:
     - name: Get occlum version
       run: echo "OCCLUM_VERSION=$(grep 'Version =' src/pal/include/occlum_version.h | awk '{print $4}')" >> $GITHUB_ENV;
 
+    - name: Get branch name
+      # GITHUB_REF=refs/heads/branch-name
+      # ${GITHUB_REF##*/} == branch-name
+      run: echo "OCCLUM_BRANCH=$(echo ${GITHUB_REF##*/})" >> $GITHUB_ENV;
+
     # Because "Build and push" step `context` field can't be subdir,
     # we need to copy files needed by dockerfile to root dir of the project
     - name: Copy context for docker build
       run: |
         cp -r tools/docker .
+
     - name: Set up QEMU
       uses: docker/setup-qemu-action@v1
 
@@ -55,7 +61,9 @@ jobs:
         context: .
         file: ./tools/docker/ci/Dockerfile.grpc
         platforms: linux/amd64
-        build-args: OCCLUM_VERSION=${{ env.OCCLUM_VERSION }}
+        build-args: |
+          "OCCLUM_VERSION=${{ env.OCCLUM_VERSION }}"
+          "OCCLUM_BRANCH=${{ env.OCCLUM_BRANCH }}"
         push: true
         tags: occlumbackup/occlum:${{ github.event.inputs.tag }}-ubuntu18.04-grpc
 
diff --git a/tools/docker/ci/Dockerfile.grpc b/tools/docker/ci/Dockerfile.grpc
index 09b75726..1c6b1228 100644
--- a/tools/docker/ci/Dockerfile.grpc
+++ b/tools/docker/ci/Dockerfile.grpc
@@ -2,9 +2,10 @@ ARG OCCLUM_VERSION
 FROM occlum/occlum:$OCCLUM_VERSION-ubuntu18.04 as base
 LABEL maintainer="Chunyang Hui <sanqian.hcy@antgroup.com>"
 
+ARG OCCLUM_BRANCH
 WORKDIR /root
 RUN rm -rf /root/demos && \
-    git clone https://github.com/occlum/occlum.git && \
+    git clone -b $OCCLUM_BRANCH https://github.com/occlum/occlum.git && \
     cp -r occlum/demos /root/demos && \
     rm -rf /root/occlum