diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml
index 6e279b1..aba4166 100644
--- a/.github/workflows/docker-image.yml
+++ b/.github/workflows/docker-image.yml
@@ -58,11 +58,15 @@ jobs:
         push: ${{ github.event_name != 'pull_request' }}
         tags:
           ${{ env.DOCKER_HUB_USERNAME }}/${{ env.DOCKER_HUB_REPOSITORY }}:pytorch1.13.1-cu116
+    - name: Build and push pytorch2.0.1-cu117 version
+      uses: docker/build-push-action@v4
+      with:
+        context: .
+        file: ./pytorch2.0.1-cu117.Dockerfile
+        push: ${{ github.event_name != 'pull_request' }}
+        tags:
+          ${{ env.DOCKER_HUB_USERNAME }}/${{ env.DOCKER_HUB_REPOSITORY }}:pytorch2.0.1-cu117
           ${{ env.DOCKER_HUB_USERNAME }}/${{ env.DOCKER_HUB_REPOSITORY }}:latest
-    # - name: Build the Docker image
-    #   run: |
-    #     docker build . --file python3.8-slim.Dockerfile --tag $DOCKER_IMAGE_TAG-python3.8-slim
-    #     docker build . --file pytorch1.13.1-cu116.Dockerfile --tag $DOCKER_IMAGE_TAG-pytorch1.13.1-cu116
     - name: Test the Docker image
       run: |
         docker images  # list local images
@@ -72,3 +76,7 @@ jobs:
         docker images  # list local images
         docker run --rm ${{ env.DOCKER_HUB_USERNAME }}/${{ env.DOCKER_HUB_REPOSITORY }}:pytorch1.13.1-cu116 bash -c "fl-sim example-configs/action-test.yml"
         docker images  # list local images
+        docker rmi ${{ env.DOCKER_HUB_USERNAME }}/${{ env.DOCKER_HUB_REPOSITORY }}:pytorch1.13.1-cu116  # remove tested image to save space
+        docker images  # list local images
+        docker run --rm ${{ env.DOCKER_HUB_USERNAME }}/${{ env.DOCKER_HUB_REPOSITORY }}:pytorch2.0.1-cu117 bash -c "fl-sim example-configs/action-test.yml"
+        docker images  # list local images
diff --git a/pytorch1.13.1-cu116.Dockerfile b/pytorch1.13.1-cu116.Dockerfile
index 2b3841a..3b9af55 100644
--- a/pytorch1.13.1-cu116.Dockerfile
+++ b/pytorch1.13.1-cu116.Dockerfile
@@ -3,7 +3,7 @@ FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime
 # NOTE:
 # pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime has python version 3.10.8, system version Ubuntu 18.04.6 LTS
 # pytorch/pytorch:1.10.1-cuda11.3-cudnn8-runtime has python version 3.7.x
-# pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime has python version 3.10.9, system version Ubuntu 18.04.6 LTS
+# pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime has python version 3.10.11, system version Ubuntu 20.04.6 LTS
 # FROM python/python:3.8-slim
 
 ## The MAINTAINER instruction sets the author field of the generated images.
diff --git a/pytorch2.0.1-cu117.Dockerfile b/pytorch2.0.1-cu117.Dockerfile
new file mode 100644
index 0000000..d78d309
--- /dev/null
+++ b/pytorch2.0.1-cu117.Dockerfile
@@ -0,0 +1,39 @@
+# https://hub.docker.com/r/pytorch/pytorch
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+# NOTE:
+# pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime has python version 3.10.8, system version Ubuntu 18.04.6 LTS
+# pytorch/pytorch:1.10.1-cuda11.3-cudnn8-runtime has python version 3.7.x
+# pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime has python version 3.10.11, system version Ubuntu 20.04.6 LTS
+# FROM python/python:3.8-slim
+
+## The MAINTAINER instruction sets the author field of the generated images.
+LABEL maintainer="wenh06@gmail.com"
+
+RUN mkdir /fl_sim
+COPY ./ /fl_sim
+WORKDIR /fl_sim
+
+
+## Install your dependencies here using apt install, etc.
+
+RUN apt update && apt upgrade -y && apt clean
+RUN apt install ffmpeg libsm6 libxext6 tar unzip wget vim nano -y
+
+# RUN apt install python3-pip
+RUN ln -s /usr/bin/python3 /usr/bin/python && ln -s /usr/bin/pip3 /usr/bin/pip
+# RUN pip install --upgrade pip
+
+# http://mirrors.aliyun.com/pypi/simple/
+# http://pypi.douban.com/simple/
+# RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
+## Include the following line if you have a requirements.txt file.
+RUN pip install -r requirements-no-torch.txt
+RUN pip install -r requirements-viz.txt
+# RUN pip install torch==2.0.1+cu117 -f https://download.pytorch.org/whl/torch_stable.html
+RUN pip install torchvision==0.15.2+cu117 --no-deps -f https://download.pytorch.org/whl/torch_stable.html
+RUN pip install torch-optimizer --no-deps
+RUN python -m pip cache purge
+
+RUN python -m pip install .
+
+# RUN python docker_test.py