Fix forward test failure.

Summary: Trying to fix this test https://www.internalfb.com/intern/test/562950062596904/

Reviewed By: tugsbayasgalan

Differential Revision: D52552898

fbshipit-source-id: fb8311f9d979ebfd3306efaac5df812766aa899b

Author: zhxchen17

exir/lowered_backend_module.py

@@ -436,7 +436,9 @@ def _get_new_signature(
                     )
                 )
         if node.op == "output":
-            for output in node.all_input_nodes:
+            for output in pytree.tree_leaves((node.args, node.kwargs)):
+                if not isinstance(output, torch.fx.Node):
+                    continue
                 output_specs.append(
                     OutputSpec(
                         kind=OutputKind.USER_OUTPUT,

exir/program/_program.py

@@ -61,9 +61,10 @@ def _get_updated_graph_signature(
     Update the graph signature's user_input/user_outputs.
     """
     new_input_specs = []
-    for i, node in enumerate(new_gm.graph.nodes):
+    i = 0
+    for node in new_gm.graph.nodes:
         if node.op != "placeholder":
-            break
+            continue
 
         assert i < len(
             old_signature.input_specs
@@ -77,6 +78,7 @@ def _get_updated_graph_signature(
         new_input_specs.append(
             InputSpec(old_input_spec.kind, arg, old_input_spec.target)
         )
+        i += 1
 
     output_node = list(new_gm.graph.nodes)[-1]
     assert output_node.op == "output"