Add support for list of leaf_module in AITTestCase (#810)

Summary: Pull Request resolved: #810 Differential Revision: D47196496 fbshipit-source-id: dbad60ab5e767877b3854f89b1c6a27d96ab7c0d
facebookincubator · Jul 4, 2023 · cecf2a8 · cecf2a8
1 parent 039bb9f
commit cecf2a8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/fx2ait/fx2ait/tools/common_fx2ait.py b/fx2ait/fx2ait/tools/common_fx2ait.py
@@ -114,7 +114,10 @@ def run_test(
 
         leaf_module_list = []
         if leaf_module:
-            leaf_module_list.append(leaf_module)
+            if isinstance(leaf_module, list):
+                leaf_module_list.extend(leaf_module)
+            else:
+                leaf_module_list.append(leaf_module)
 
         orig_mod = copy.deepcopy(mod)
         orig_mod.eval()