Update on "[standalone_compile] Fix single Tensor outputs from split_module"

zou3519 · zou3519 · commit 281ca0c84272 · 2025-07-08T09:33:51.000-07:00
We assumed that the output in an FX graph would always just be a
list[Tensor], even in the single tensor return case.
It is possible for the output to be a single Tensor. This can happen
by calling torch.fx.split_module on the module.

Test Plan:
- new test

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov

[ghstack-poisoned]
diff --git a/torch/_inductor/standalone_compile.py b/torch/_inductor/standalone_compile.py
@@ -214,7 +214,7 @@ def standalone_compile(
         assert last_node.op == "output"
         assert len(last_node.args) == 1
 
-        def handle_node(node):
+        def handle_node(node: torch.fx.Node) -> None:
             nonlocal fake_mode
             if "example_value" in node.meta:
                 maybe_tensor = node.meta["example_value"]
