From dc236b4f72c52de1b1fac8106bb9e1aca801ae3e Mon Sep 17 00:00:00 2001
From: Mark Saroufim <marksaroufim@meta.com>
Date: Mon, 11 Nov 2024 16:23:50 -0800
Subject: [PATCH] update

---
 .github/workflows/train_workflow.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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