Fix github action CI for GPU (#625)

Summary: Pull Request resolved: #625

Differential Revision: D53292685

.github/workflows/ci_gpu.yml

@@ -8,16 +8,17 @@ on:
     branches:
       - main
 
+jobs:
   unittest_multi_gpu:
     runs-on: linux.4xlarge.nvidia.gpu
     steps:
       - name: Checkout
         uses: actions/checkout@v2
 
-      - name: Set up Python
-        uses: actions/setup-python@v2
-        with:
-          python-version: 3.9
+      # - name: Set up Python
+      #   uses: actions/setup-python@v2
+      #   with:
+      #     python-version: '3.9'
 
       - name: Install dependencies
         run: |
@@ -31,21 +32,21 @@ on:
 
 
   integrationtest_py39_torch_release_cuda:
-    runs-on: ubuntu-latest
-    container:
-      # https://hub.docker.com/r/nvidia/cuda
-      image: nvidia/cuda:12.3.1-base-ubuntu22.04
-      options: --gpus all
-    env:
-      TZ: 'UTC'
+    runs-on: linux.4xlarge.nvidia.gpu
+    # container:
+    #   # https://hub.docker.com/r/nvidia/cuda
+    #   image: nvidia/cuda:12.3.1-base-ubuntu22.04
+    #   options: --gpus all
+    # env:
+    #   TZ: 'UTC'
     steps:
       - name: Checkout
         uses: actions/checkout@v2
 
-      - name: Set up Python
-        uses: actions/setup-python@v2
-        with:
-          python-version: 3.9
+      # - name: Set up Python
+      #   uses: actions/setup-python@v2
+      #   with:
+      #     python-version: 3.9
 
       - name: Install dependencies
         run: |