Flatten output arguments before returning them

cchung100m · web-flow · commit ccd4bbd1e824 · 2026-05-02T18:39:21.000+08:00
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1338,7 +1338,7 @@ def _translate_fx_graph(
                 raise ValueError(f"Unsupported op {node.op}")
 
         assert output_args is not None
-        return output_args
+        return self._flatten_output_args(output_args)
 
     @staticmethod
     def _flatten_output_args(output_args) -> tuple[relax.Expr, ...]:
