OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[graph_manipulation] Unpack list of outputs (#72940)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72940

att

Reviewed By: jackm321

Differential Revision: D34282062

fbshipit-source-id: 743710c18e1f38286d1b91c91868bb22c760f3ca
This commit is contained in:
Jordan Fix 2022-02-16 22:31:51 -08:00 committed by Facebook GitHub Bot
parent fb9e92fea5
commit fd2bdd189d
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
torch/fx/passes/graph_manipulation.py
View File

@ -451,9 +451,9 @@ def serialize_module(fx_module: GraphModule, weights: Dict, name_prefix="") -> D
get_output_arg_info,
)
# If there're multiple outputs then node_rep["args"][0] will be a tuple.
# In this case we want to unpack the tuple.
if isinstance(node_rep["args"][0], tuple):
# If there're multiple outputs then node_rep["args"][0] will be a tuple or
# list. In this case we want to unpack the tuple or list.
if isinstance(node_rep["args"][0], (tuple, list)):
node_rep["args"] = node_rep["args"][0]
else:
node_rep["args"] = map_aggregate(node.args, get_arg_info)