Fix failing tests (#857)

Summary: Pull Request resolved: #857 Reviewed By: frank-wei, mortzur Differential Revision: D47874377 fbshipit-source-id: 8afe4e107e9c7ba16aed3e8931de3e64d7cb11c1
facebookincubator · Jul 28, 2023 · 1440635 · 1440635
1 parent ae72a43
commit 1440635
Showing 1 changed file with 14 additions and 10 deletions.
diff --git a/fx2ait/fx2ait/tools/common_aten2ait.py b/fx2ait/fx2ait/tools/common_aten2ait.py
@@ -102,16 +102,20 @@ def generate_graph(
         if customized_passes:
             passes_list.extend(customized_passes)
 
-        fx_module = exir.capture(
-            mod,
-            tuple(original_inputs),
-            CaptureConfig(
-                pt2_mode=True,
-                enable_functionalization=False,
-                enable_dynamic_shape=True,
-                _use_old_decomp_table=True,
-            ),
-        ).transform(*tuple(passes_list))
+        fx_module = (
+            exir.capture(
+                mod,
+                tuple(original_inputs),
+                CaptureConfig(
+                    pt2_mode=True,
+                    enable_functionalization=False,
+                    enable_dynamic_shape=True,
+                    _use_old_decomp_table=True,
+                ),
+            )
+            .transform(*tuple(passes_list))
+            .exported_program.graph_module
+        )
 
         fx_module = run_const_fold(fx_module)
         _LOGGER.info(f"aten fx graph: {fx_module.graph}")