From e17a43184b6e28afb6a4757076d27018cbf8c5d0 Mon Sep 17 00:00:00 2001
From: Frank Lee <somerlee.9@gmail.com>
Date: Tue, 24 May 2022 16:53:39 +0800
Subject: [PATCH] [ci] update the docker image name (#1017)

---
 .github/workflows/build.yml         | 2 +-
 .github/workflows/release_bdist.yml | 4 ++--
 docker/Dockerfile                   | 7 +------
 3 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index dbe7315b1..bcece5ab3 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -14,7 +14,7 @@ jobs:
         contains( github.event.pull_request.labels.*.name, 'Run Build and Test')
     runs-on: [self-hosted, gpu]
     container:
-      image: frankleeeee/pytorch-cuda:1.10.1-11.3.0
+      image: hpcaitech/pytorch-cuda:1.10.1-11.3.0
       options: --gpus all --rm -v /data/scratch/cifar-10:/data/scratch/cifar-10
     timeout-minutes: 40
     steps:
diff --git a/.github/workflows/release_bdist.yml b/.github/workflows/release_bdist.yml
index ab60ecc63..242d02d4b 100644
--- a/.github/workflows/release_bdist.yml
+++ b/.github/workflows/release_bdist.yml
@@ -28,9 +28,9 @@ jobs:
     steps:
     - id: set-matrix
       run: |
-        [ "${{github.event.inputs.cuda_version}}" != "" ] && matrix="[\"frankleeeee/cuda-conda:${{github.event.inputs.cuda_version}}\"]"
+        [ "${{github.event.inputs.cuda_version}}" != "" ] && matrix="[\"hpcaitech/cuda-conda:${{github.event.inputs.cuda_version}}\"]"
         [ "${{github.event.inputs.cuda_version}}" == "" ] || [ "${{github.event.inputs.version}}" == "all" ] && \
-          matrix="[\"frankleeeee/cuda-conda:11.3\", \"frankleeeee/cuda-conda:11.1\", \"frankleeeee/cuda-conda:10.2\"]"
+          matrix="[\"hpcaitech/cuda-conda:11.3\", \"hpcaitech/cuda-conda:11.1\", \"hpcaitech/cuda-conda:10.2\"]"
         echo $matrix
         echo "::set-output name=matrix::{\"container\":$(echo $matrix)}"
 
diff --git a/docker/Dockerfile b/docker/Dockerfile
index da851a156..123c7235b 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,9 +1,4 @@
-FROM nvcr.io/nvidia/pytorch:21.07-py3
-
-# install dependencies
-RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple \
- && pip install -U pip setuptools wheel \
- && pip install pytest tensorboard deepspeed apex
+FROM hpcaitech/pytorch-cuda:1.10.1-11.3.0
 
 # install colossalai
 RUN git clone https://github.com/hpcaitech/ColossalAI.git \