From cecf2a899e86e46144214b834fa21a24a8b888d7 Mon Sep 17 00:00:00 2001 From: Henry Hu Date: Mon, 3 Jul 2023 19:35:32 -0700 Subject: [PATCH] Add support for list of leaf_module in AITTestCase (#810) Summary: Pull Request resolved: https://github.com/facebookincubator/AITemplate/pull/810 Differential Revision: D47196496 fbshipit-source-id: dbad60ab5e767877b3854f89b1c6a27d96ab7c0d --- fx2ait/fx2ait/tools/common_fx2ait.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/fx2ait/fx2ait/tools/common_fx2ait.py b/fx2ait/fx2ait/tools/common_fx2ait.py index 898a14234..8d3f85964 100644 --- 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()