Fix github action CI for GPU (#625)

Summary: Pull Request resolved: #625

Differential Revision: D53292685

.github/workflows/ci_gpu.yml

@@ -8,44 +8,75 @@ 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: Install dependencies
+  #       run: |
+  #         ./scripts/install_via_pip.sh -c
+
+  #     - name: Run multi-GPU unit tests
+  #       run: |
+  #         nvidia-smi
+  #         nvcc --version
+  #         python -m unittest opacus.tests.multigpu_gradcheck.GradientComputationTest.test_gradient_correct
   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: Install dependencies
-        run: |
-          ./scripts/install_via_pip.sh -c
-
-      - name: Run multi-GPU unit tests
-        run: |
-          nvidia-smi
-          nvcc --version
-          python -m unittest opacus.tests.multigpu_gradcheck.GradientComputationTest.test_gradient_correct
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      runner: linux.4xlarge.nvidia.gpu
+      gpu-arch-type: cuda
+      gpu-arch-version: "11.6"
+      script: |
+        ./scripts/install_via_pip.sh -c
+        nvidia-smi
+        nvcc --version
+        python -m unittest opacus.tests.multigpu_gradcheck.GradientComputationTest.test_gradient_correct
+      # 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: Install dependencies
+      #     run: |
+      #       ./scripts/install_via_pip.sh -c
+
+      #   - name: Run multi-GPU unit tests
+      #     run: |
+      #       nvidia-smi
+      #       nvcc --version
+      #       python -m unittest opacus.tests.multigpu_gradcheck.GradientComputationTest.test_gradient_correct
 
 
   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
+          python-version: '3.10.13'
 
       - name: Install dependencies
         run: |