Automatically enable GPU in train script when gpu flag is provided

train.sh

@@ -63,6 +63,7 @@ fi
 
 if [ ! -z "$gpu" ]; then
     docker_params+=(--gpus device=$gpu)
+    script_params+=(--num_gpus 1)
 fi
 
 mkdir -p $output_dir