Update 3.container-train.sbatch

aws-samples · Mar 25, 2024 · c2ea6b7 · c2ea6b7
1 parent b1b4e5f
commit c2ea6b7
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/3.test_cases/16.pytorch-cpu-ddp/3.container-train.sbatch b/3.test_cases/16.pytorch-cpu-ddp/3.container-train.sbatch
@@ -17,6 +17,7 @@ head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)
 
 echo Node IP: $head_node_ip
 export LOGLEVEL=INFO
+export NVIDIA_VISIBLE_DEVICES=void
 
 declare -a ARGS=(
     --container-image ${PWD}/pytorch.sqsh