diff --git a/.github/workflows/release_docker_after_merge.yml b/.github/workflows/release_docker_after_merge.yml
index 607c19b05..6b6684641 100644
--- a/.github/workflows/release_docker_after_merge.yml
+++ b/.github/workflows/release_docker_after_merge.yml
@@ -26,7 +26,7 @@ jobs:
         run: |
           version=$(cat version.txt)
           tag=hpcaitech/colossalai:$version
-          docker build --build-arg http_proxy=http://172.17.0.1:7890 --build-arg https_proxy=http://172.17.0.1:7890 -t $tag ./docker
+	  docker build --build-arg http_proxy=http://172.17.0.1:7890 --build-arg https_proxy=http://172.17.0.1:7890 --build-arg VERSION=v${version} -t $tag ./docker
           echo "tag=${tag}" >> $GITHUB_OUTPUT
 
       - name: Log in to Docker Hub
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 49ff9b344..52c7bf560 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -15,7 +15,8 @@ RUN git clone https://github.com/NVIDIA/apex && \
     pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_layer_norm" ./
 
 # install colossalai
-RUN git clone https://github.com/hpcaitech/ColossalAI.git \
+ARG VERSION=1
+RUN git clone -b ${VERSION} https://github.com/hpcaitech/ColossalAI.git \
     && cd ./ColossalAI \
     && CUDA_EXT=1 pip install -v --no-cache-dir .