Fix handling of extra_args for ModelBuilder pass (microsoft#1118)

DavitGrigoryan132 · Aug 14, 2024 · 6f7aba7 · 6f7aba7
1 parent bc6b959
commit 6f7aba7
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/olive/passes/onnx/model_builder.py b/olive/passes/onnx/model_builder.py
@@ -186,9 +186,14 @@ def _run_for_config(
         if config.get("int4_accuracy_level"):
             extra_args["int4_accuracy_level"] = config["int4_accuracy_level"].value
 
-        extra_args["exclude_embeds"] = config["exclude_embeds"]
-        extra_args["exclude_lm_head"] = config["exclude_lm_head"]
-        extra_args["enable_cuda_graph"] = "1" if config["enable_cuda_graph"] else "0"
+        if config["exclude_embeds"]:
+            extra_args["exclude_embeds"] = config["exclude_embeds"]
+
+        if config["exclude_lm_head"]:
+            extra_args["exclude_lm_head"] = config["exclude_lm_head"]
+
+        if config["enable_cuda_graph"] is not None:
+            extra_args["enable_cuda_graph"] = "1" if config["enable_cuda_graph"] else "0"
 
         create_model(
             model_name=model_path,