update

gpu-mode · Nov 12, 2024 · dc236b4 · dc236b4
1 parent 66495dd
commit dc236b4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/.github/workflows/train_workflow.yml b/.github/workflows/train_workflow.yml
@@ -11,7 +11,8 @@ on:
 jobs:
   train:
     runs-on: [gpumode-nvidia-arc]
-
+    container:
+      image: nvidia/cuda:12.4.0-devel-ubuntu22.04
     steps:
     - name: Setup Python
       uses: actions/setup-python@v4