diff --git a/.github/scripts/docker/Dockerfile b/.github/scripts/docker/Dockerfile
index 03e598466..47af9983b 100644
--- a/.github/scripts/docker/Dockerfile
+++ b/.github/scripts/docker/Dockerfile
@@ -1,7 +1,15 @@
 ARG PYTHON_VERSION=3.8
-
 FROM python:${PYTHON_VERSION}
 
+ARG TORCHAUDIO_VERSION="0.13.0"
+ARG TORCH_VERSION="1.13.0"
+ARG K2_VERSION="1.24.4.dev20231220"
+ARG KALDIFEAT_VERSION="1.25.3.dev20231221"
+
+ARG _K2_VERSION="${K2_VERSION}+cpu.torch${TORCH_VERSION}"
+ARG _KALDIFEAT_VERSION="${KALDIFEAT_VERSION}+cpu.torch${TORCH_VERSION}"
+
+
 RUN apt-get update -y && \
     apt-get install -qq -y \
     git \
@@ -12,22 +20,18 @@ RUN apt-get update -y && \
     apt-get clean && \
     rm -rf /var/cache/apt/archives /var/lib/apt/lists
 
-ARG K2_VERSION="1.24.4.dev20231220+cpu.torch1.13.0"
-ARG KALDIFEAT_VERSION="1.25.3.dev20231221+cpu.torch1.13.0"
-ARG TORCHAUDIO_VERSION="0.13.0"
-ARG TORCH_VERSION="1.13.0"
 
 LABEL authors="Fangjun Kuang <csukuangfj@gmail.com>"
-LABEL k2_version=${K2_VERSION}
-LABEL kaldifeat_version=${KALDIFEAT_VERSION}
+LABEL k2_version=${_K2_VERSION}
+LABEL kaldifeat_version=${_KALDIFEAT_VERSION}
 LABEL github_repo="https://github.com/k2-fsa/icefall"
 
 # Install dependencies
 RUN pip install --no-cache-dir \
       torch==${TORCH_VERSION} torchaudio==${TORCHAUDIO_VERSION} -f https://download.pytorch.org/whl/cpu/torch_stable.html \
-      k2==${K2_VERSION} -f https://k2-fsa.github.io/k2/cpu.html \
+      k2==${_K2_VERSION} -f https://k2-fsa.github.io/k2/cpu.html \
       git+https://github.com/lhotse-speech/lhotse \
-      kaldifeat==${KALDIFEAT_VERSION} -f https://csukuangfj.github.io/kaldifeat/cpu.html \
+      kaldifeat==${_KALDIFEAT_VERSION} -f https://csukuangfj.github.io/kaldifeat/cpu.html \
       kaldi_native_io \
       kaldialign \
       kaldifst \
diff --git a/.github/workflows/build-cpu-docker.yml b/.github/workflows/build-cpu-docker.yml
index 67af26397..bf1125779 100644
--- a/.github/workflows/build-cpu-docker.yml
+++ b/.github/workflows/build-cpu-docker.yml
@@ -2,7 +2,7 @@ name: build-cpu-docker
 on:
   push:
     branches:
-      - ci-train-2
+      - ci-train-3
   workflow_dispatch:
 
 concurrency:
@@ -11,13 +11,16 @@ concurrency:
 
 jobs:
   build-cpu-docker:
-    name: ${{ matrix.python-version }}
+    name: py${{ matrix.python-version }} torch${{ matrix.torch-version }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         os: [ubuntu-latest]
         python-version: ["3.8", "3.9", "3.10"]
+        torch-version: ["1.13.0", "2.0.0", "2.0.1", "2.1.0", "2.1.1", "2.1.2"]
+        k2-version: ["1.24.4.dev20231220"]
+        kaldifeat-version: ["1.25.3.dev20231221"]
 
     steps:
       # refer to https://github.com/actions/checkout
@@ -43,7 +46,30 @@ jobs:
         shell: bash
         run: |
           cd .github/scripts/docker
+          torch_version=${{ matrix.torch-version }}
+          if [[ $torch_version == 1.13.0 ]]; then
+            torchaudio_version=0.13.0
+          elif [[ $torch_version == 2.0.0 ]]; then
+            torchaudio_version=2.0.1
+          elif [[ $torch_version == 2.0.1 ]]; then
+            torchaudio_version=2.0.2
+          else
+            torchaudio_version=$torch_version
+          fi
+          echo "torch_version: $torch_version"
+          echo "torchaudio_version: $torchaudio_version"
+
+          tag=ghcr.io/csukuangfj/icefall:cpu-py${{ matrix.python-version }}-torch${{ matrix.torch-version }}
+          echo "tag: $tag"
+
+          docker build \
+            -t $tag \
+            --build-arg PYTHON_VERSION=${{ matrix.python-version }} \
+            --build-arg TORCH_VERSION=$torch_version \
+            --build-arg TORCHAUDIO_VERSION=$torchaudio_version \
+            --build-arg K2_VERSION=${{ matrix.k2-version }} \
+            --build-arg KALDIFEAT_VERSION=${{ matrix.kaldifeat-version }} \
+            .
 
-          docker build -t ghcr.io/csukuangfj/icefall:cpu-py${{ matrix.python-version }} --build-arg PYTHON_VERSION=${{ matrix.python-version }} .
           docker image ls
-          docker push ghcr.io/csukuangfj/icefall:cpu-py${{ matrix.python-version }}
+          docker push $tag
diff --git a/.github/workflows/train-librispeech.yml b/.github/workflows/train-librispeech.yml
index 1e7ee8bc3..0a76e0cd7 100644
--- a/.github/workflows/train-librispeech.yml
+++ b/.github/workflows/train-librispeech.yml
@@ -2,7 +2,13 @@ name: train librispeech
 on:
   push:
     branches:
-      - ci-train
+      - master
+      - ci-train-2
+
+  pull_request:
+    branches:
+      - master
+
   workflow_dispatch:
 
 concurrency:
@@ -11,13 +17,16 @@ concurrency:
 
 jobs:
   train-librispeech:
-    name: ${{ matrix.python-version }}
+    name: py${{ matrix.python-version }} torch${{ matrix.torch-version }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         os: [ubuntu-latest]
         python-version: ["3.8", "3.9", "3.10"]
+        torch-version: ["1.13.0", "2.0.0", "2.0.1", "2.1.0", "2.1.1", "2.1.2"]
+        k2-version: ["1.24.4.dev20231220"]
+        kaldifeat-version: ["1.25.3.dev20231221"]
 
     steps:
       # refer to https://github.com/actions/checkout
@@ -37,7 +46,8 @@ jobs:
       - name: Run the build process with Docker
         uses: addnab/docker-run-action@v3
         with:
-            image: ghcr.io/csukuangfj/icefall:cpu-py${{ matrix.python-version }}
+            # image: ghcr.io/k2-fsa/icefall:cpu-py${{ matrix.python-version }}
+            image: ghcr.io/csukuangfj/icefall:cpu-py${{ matrix.python-version }}-torch${{ matrix.torch-version }}
             options: |
               --volume ${{ github.workspace }}/:/icefall
             shell: bash