gpus all

huggingface · Sep 14, 2024 · 87c9f3e · 87c9f3e
1 parent 90aa85d
commit 87c9f3e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/.github/workflows/test_onnxruntime_gpu.yml b/.github/workflows/test_onnxruntime_gpu.yml
@@ -18,6 +18,7 @@ jobs:
 
     container:
       image: nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
+      options: --gpus all
 
     steps:
       - name: Checkout

diff --git a/.github/workflows/test_onnxruntime_train.yml b/.github/workflows/test_onnxruntime_train.yml
@@ -13,7 +13,8 @@ jobs:
       group: aws-g6-4xlarge-plus
 
     container:
-      image: nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
+      image: nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
+      options: --gpus all
 
     steps:
       - name: Checkout