diff --git a/atlas_server/src/static/training/predict.job.sh b/atlas_server/src/static/training/predict.job.sh
index 1ec14ebad696b0f91bd502c3791ae8fa61e480d6..c3890159adf4b6fc9ea4f5d5048cd64933510d48 100644
--- a/atlas_server/src/static/training/predict.job.sh
+++ b/atlas_server/src/static/training/predict.job.sh
@@ -17,6 +17,8 @@ module use ${SOFTWARE_STAGE_MODULES} &> /dev/null
 module load ${ATLAS_MODULE}
 echo Modules loaded
 
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
 # Environment variables
 export SCRATCH_USER=/p/scratch/cjinm16/schiffer1/
 export OUTPUT_SPACING=16
diff --git a/atlas_server/src/static/training/train.job.sh b/atlas_server/src/static/training/train.job.sh
index b5651e9a49f731ee3fbb27192ed0bc2e9ea7b5b1..d0249f9769441e1c3fa1922134aeddede9c687c9 100644
--- a/atlas_server/src/static/training/train.job.sh
+++ b/atlas_server/src/static/training/train.job.sh
@@ -17,6 +17,8 @@ module use ${SOFTWARE_STAGE_MODULES} &> /dev/null
 module load ${ATLAS_MODULE}
 echo Modules loaded
 
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
 # Environment variables
 export BRAIN=__BRAIN__
 export LABELS=__LABELS__