From 8da68ee0b5da4911a1d632c6babc9c5bdb765a97 Mon Sep 17 00:00:00 2001 From: panfengfeng Date: Wed, 23 Sep 2020 14:12:42 +0800 Subject: [PATCH] fix deepfm gpu train script --- .../deepfm/scripts/run_standalone_train.sh | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/model_zoo/official/recommend/deepfm/scripts/run_standalone_train.sh b/model_zoo/official/recommend/deepfm/scripts/run_standalone_train.sh index f2c5c16eed3..7780500b11c 100644 --- a/model_zoo/official/recommend/deepfm/scripts/run_standalone_train.sh +++ b/model_zoo/official/recommend/deepfm/scripts/run_standalone_train.sh @@ -14,12 +14,22 @@ # limitations under the License. # ============================================================================ echo "Please run the script as: " -echo "sh scripts/run_standalone_train.sh DEVICE_ID DEVICE_TARGET DATASET_PATH" +echo "sh scripts/run_standalone_train.sh DEVICE_ID/CUDA_VISIBLE_DEVICES DEVICE_TARGET DATASET_PATH" echo "for example: sh scripts/run_standalone_train.sh 0 GPU /dataset_path" echo "After running the script, the network runs in the background, The log will be generated in ms_log/output.log" -export DEVICE_ID=$1 DEVICE_TARGET=$2 + +if [ "$DEVICE_TARGET" = "GPU" ] +then + export CUDA_VISIBLE_DEVICES=$1 +fi + +if [ "$DEVICE_TARGET" = "Ascend" ]; +then + export DEVICE_ID=$1 +fi + DATA_URL=$3 mkdir -p ms_log