Fix boolean parsing for capture-onnx-graph options

olive/cli/capture_onnx.py

@@ -147,23 +147,18 @@ def register_subcommand(parser: ArgumentParser):
         )
         mb_group.add_argument(
             "--exclude_embeds",
-            type=bool,
-            default=False,
-            required=False,
+            action="store_true",
             help="Remove embedding layer from your ONNX model.",
         )
         mb_group.add_argument(
             "--exclude_lm_head",
-            type=bool,
-            default=False,
-            required=False,
+            action="store_true",
             help="Remove language modeling head from your ONNX model.",
         )
         mb_group.add_argument(
             "--enable_cuda_graph",
-            type=bool,
+            action="store_true",
             default=None,  # Explicitly setting to None to differentiate between user intent and default.
-            required=False,
             help=(
                 "The model can use CUDA graph capture for CUDA execution provider. "
                 "If enabled, all nodes being placed on the CUDA EP is the prerequisite "

test/cli/test_capture_onnx_args.py

@@ -0,0 +1,34 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import argparse
+
+from olive.cli.capture_onnx import CaptureOnnxGraphCommand
+
+
+def _parse_capture_args(*args):
+    parser = argparse.ArgumentParser()
+    commands = parser.add_subparsers()
+    CaptureOnnxGraphCommand.register_subcommand(commands)
+    return parser.parse_args(["capture-onnx-graph", *args])
+
+
+def test_capture_onnx_boolean_flags_default_to_disabled():
+    args = _parse_capture_args()
+
+    assert args.exclude_embeds is False
+    assert args.exclude_lm_head is False
+    assert args.enable_cuda_graph is None
+
+
+def test_capture_onnx_boolean_flags_enable_on_presence():
+    args = _parse_capture_args(
+        "--exclude_embeds",
+        "--exclude_lm_head",
+        "--enable_cuda_graph",
+    )
+
+    assert args.exclude_embeds is True
+    assert args.exclude_lm_head is True
+    assert args.enable_cuda_graph is True