Update frontend.py

buddy-compiler · Nov 18, 2024 · b380973 · b380973
1 parent f45d990
commit b380973
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/frontend/Python/frontend.py b/frontend/Python/frontend.py
@@ -331,8 +331,8 @@ def _compiler(_gm: torch.fx.GraphModule, _inputs: List[torch.Tensor]):
                 else:
                     tensor_meta = gm_node.meta.get("tensor_meta")
                     val = gm_node.meta.get("val")
-                    num_returns = len(gm_node.target._schema.returns)
-                    # num_returns = len(val) if isinstance(val, list) else len(gm_node.target._schema.returns)
+                    # num_returns = len(gm_node.target._schema.returns)
+                    num_returns = len(val) if isinstance(val, list) else len(gm_node.target._schema.returns)
                     if num_returns == 1:
                         node_dtype = self._torch_dtype_translate(
                             str(tensor_meta.dtype)