diff --git a/.github/workflows/train_workflow.yml b/.github/workflows/train_workflow.yml
index 4c87fc1..008440a 100644
--- 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