diff --git a/tools/mace_tools.py b/tools/mace_tools.py
index 27908389d8c8c50a44d8de6c74b72914a90beea4..dc43b1cec981756246e30546a4598a108571cb93 100644
--- a/tools/mace_tools.py
+++ b/tools/mace_tools.py
@@ -369,7 +369,8 @@ def process_models(project_name, configs, embed_model_data, vlog_level,
                                          model_config["input_shapes"],
                                          input_file_list)
 
-        if FLAGS.mode == "build" or FLAGS.mode == "all":
+        if FLAGS.mode == "build" or FLAGS.mode == "benchmark" or \
+                FLAGS.mode == "all":
             sh_commands.gen_model_code(
                     "mace/codegen/models/%s" % model_name,
                     model_config["platform"],