diff --git a/.github/workflows/compatibility_test.yml b/.github/workflows/compatibility_test.yml
index 0e3eddb9b..e41ae9379 100644
--- a/.github/workflows/compatibility_test.yml
+++ b/.github/workflows/compatibility_test.yml
@@ -6,14 +6,14 @@ jobs:
   build:
     name: Test for PyTorch compatibility
     if: github.base_ref == 'main' && github.repository == 'hpcaitech/ColossalAI'
-    runs-on: [self-hosted, aws]
+    runs-on: [self-hosted, gpu]
     strategy:
       fail-fast: false
       matrix:
-        # 20.07: PyTorch 1.6.0a0+9907a3e + Python 3.6
-        # 20.10: PyTorch 1.7.0a0+7036e91 + Python 3.8
-        # 20.12: PyTorch 1.8.0a0+1606899 + Python 3.8
-        # 21.06: PyTorch 1.9.0a0+c3d40fd + Python 3.8
+        # pytorch-cuda:1.9.0-11.1.1 + Python 3.8
+        # pytorch-cuda:1.8.1-11.1.1 + Python 3.8
+        # pytorch-cuda:1.7.1-11.0.3 + Python 3.8
+        # pytorch-cuda:1.6.0-10.2 + Python 3.6
         container: ["frankleeeee/pytorch-cuda:1.9.0-11.1.1",
                     "frankleeeee/pytorch-cuda:1.8.1-11.1.1",
                     "frankleeeee/pytorch-cuda:1.7.1-11.0.3",
@@ -27,7 +27,7 @@ jobs:
         run: |
           pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
           pip install -U pip setuptools wheel --user
-          pip install pytest tensorboard deepspeed apex
+          pip install pytest tensorboard
       - uses: actions/checkout@v2
       - name: Install Colossal-AI      
         run: |