allow multi-gpu (#31011)

* allow multi-gpu

* allow multi-gpu

---------
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

.github/workflows/ssh-runner.yml

@@ -9,6 +9,9 @@ on:
       docker_image:
         description: 'Name of the Docker image'
         required: true
+      num_gpus:
+        description: 'Type of the number of gpus to use (`single` or `multi`)'
+        required: true
 
 env:
   IS_GITHUB_CI: "1"
@@ -20,12 +23,13 @@ env:
   RUN_SLOW: yes # For gated repositories, we still need to agree to share information on the Hub repo. page in order to get access. # This token is created under the bot `hf-transformers-bot`. 
   SIGOPT_API_TOKEN: ${{ secrets.SIGOPT_API_TOKEN }} 
   TF_FORCE_GPU_ALLOW_GROWTH: true 
+  CUDA_VISIBLE_DEVICES: 0,1
   RUN_PT_TF_CROSS_TESTS: 1
 
 jobs:
   ssh_runner:
     name: "SSH"
-    runs-on: [single-gpu, nvidia-gpu, "${{ github.event.inputs.runner_type }}", ci]
+    runs-on: ["${{ github.event.inputs.num_gpus }}-gpu", nvidia-gpu, "${{ github.event.inputs.runner_type }}", ci]
     container:
       image: ${{ github.event.inputs.docker_image }}
       options: --gpus all --privileged --ipc host -v /mnt/cache/.cache/huggingface:/mnt/cache/