[single_stage_detector] Updated run_and_time.sh for customizing number of GPUs on single node (#808)

* Added support for dataset customization

* Added support for setting GPUs on single node training

* Revert "Added support for dataset customization"

This reverts commit c37ae36.

Author: jamesljlster

single_stage_detector/ssd/run_and_time.sh

@@ -38,6 +38,7 @@ NUMEPOCHS=${NUMEPOCHS:-30}
 LOG_INTERVAL=${LOG_INTERVAL:-20}
 DATASET_DIR=${DATASET_DIR:-"/datasets/open-images-v6-mlperf"}
 TORCH_HOME=${TORCH_HOME:-"$(pwd)/torch-model-cache"}
+DGXNGPU=${DGXNGPU:-1}
 
 # Handle MLCube parameters
 while [ $# -gt 0 ]; do
@@ -76,7 +77,7 @@ if [ -n "${SLURM_LOCALID-}" ]; then
   fi
 else
   # Mode 2: Single-node Docker; need to launch tasks with torchrun
-  CMD=( "torchrun" "--standalone" "--nnodes=1" "--nproc_per_node=1" )
+  CMD=( "torchrun" "--standalone" "--nnodes=1" "--nproc_per_node=${DGXNGPU}" )
   [ "$MEMBIND" = false ] &&  CMD+=( "--no_membind" )
 fi