diff --git a/benchmark.sh b/benchmark.sh
index 7e40464281cbf82a227c7ec59ee3dd6f602cc7c4..2bf14a6be8b611a4e7a6a2ef17903a0d83d70c05 100644
--- a/benchmark.sh
+++ b/benchmark.sh
@@ -22,7 +22,7 @@ if [ "$EMBED_MODEL_DATA" = 0 ]; then
   cp codegen/models/${MODEL_TAG}/${MODEL_TAG}.data $MODEL_OUTPUT_DIR
 fi
 
-if [ x"$RUNTIME" = x"host" ]; then
+if [ x"$TARGET_ABI" == x"host" ]; then
   bazel build --verbose_failures -c opt --strip always benchmark:benchmark_model \
     --copt="-std=c++11" \
     --copt="-D_GLIBCXX_USE_C99_MATH_TR1" \
@@ -30,15 +30,17 @@ if [ x"$RUNTIME" = x"host" ]; then
     --copt="-DMACE_MODEL_TAG=${MODEL_TAG}" \
     --copt="-O3" \
     --define openmp=true \
-    --model_data_file=${MODEL_OUTPUT_DIR}/${MODEL_TAG}.data \
     --define production=true || exit 1
 
   cp bazel-bin/benchmark/benchmark_model $MODEL_OUTPUT_DIR
 
   MACE_CPP_MIN_VLOG_LEVEL=$VLOG_LEVEL \
   ${MODEL_OUTPUT_DIR}/benchmark_model \
+      --model_data_file=${PHONE_DATA_DIR}/${MODEL_TAG}.data \
       --device=${DEVICE_TYPE} \
+      --input_node="${INPUT_NODE}" \
       --input_shape="${INPUT_SHAPE}"\
+      --output_node="${OUTPUT_NODE}" \
       --output_shape="${OUTPUT_SHAPE}"\
       --input_file=${MODEL_OUTPUT_DIR}/${INPUT_FILE_NAME} || exit 1
 
diff --git a/mace_tools.py b/mace_tools.py
index c3e0069d143ee10d4de443fe3888c20e41310abf..905c9792281c72cd9ef475ae4c621e58e5d55214 100644
--- a/mace_tools.py
+++ b/mace_tools.py
@@ -235,7 +235,8 @@ def main(unused_args):
         os.environ["WEIGHT_FILE_PATH"] = model_output_dir + "/model.caffemodel"
         urllib.urlretrieve(model_config["weight_file_path"], os.environ["WEIGHT_FILE_PATH"])
 
-      if FLAGS.mode == "build" or FLAGS.mode == "run" or FLAGS.mode == "validate" or FLAGS.mode == "all":
+      if FLAGS.mode == "build" or FLAGS.mode == "run" or FLAGS.mode == "validate"\
+          or FLAGS.mode == "benchmark" or FLAGS.mode == "all":
         generate_random_input(model_output_dir)
 
       if FLAGS.mode == "build" or FLAGS.mode == "all":