CustomTracer.module_leaves was silently ignored because is_leaf_module had a return before the module_leaves check, making the entire feature non-functional. Additionally, the tracing path in ExportOptions.export had leftover debug artifacts.

Changes

• tracing.py: is_leaf_module now correctly delegates to module_leaves after checking standard torch.nn leaf status (dead return removed in prior commit 77a1c36)
• export_options.py:
  • Removed unconditional print(graph) that fired on every traced export
  • Fixed typo [ExportIptions.export] → [ExportOptions.export]
  • Removed commented-out dead code
• test_tracing.py: Removed unused is_leaf_module local function from test_export_submodule; the working pattern is now demonstrated in test_export_and_trace_submodule

Usage

# Use CustomTracer with module_leaves to keep selected submodules as call_module nodes
module_leaves = {SubModule: lambda mod, module_qualified_name=None: mod.kind == 1}

# Via CustomTracer directly
graph = CustomTracer(module_leaves=module_leaves).trace(model)

# Or via to_onnx with tracing
onx = to_onnx(
    model,
    (x, y),
    export_options=ExportOptions(tracing=True, tracing_module_leaves=module_leaves),
    export_modules_as_functions={"suba"},
    inline=False,
)

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.