From b7c9e2e19acb9efffcef1f5cb3da48895ff1ab05 Mon Sep 17 00:00:00 2001 From: dlyakhov Date: Mon, 25 Nov 2024 15:57:35 +0100 Subject: [PATCH] Code comments --- nncf/experimental/torch/fx/model_transformer.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/nncf/experimental/torch/fx/model_transformer.py b/nncf/experimental/torch/fx/model_transformer.py index 404b708f366..2be4fbde877 100644 --- a/nncf/experimental/torch/fx/model_transformer.py +++ b/nncf/experimental/torch/fx/model_transformer.py @@ -163,14 +163,18 @@ def remap_fn(node: torch.fx.Node): for idx, node in enumerate(nodes_with_output): if isinstance(node, torch.fx.Node): continue + # Current node is the original graph output. + # Should be replaced by its arguments. output_node = get_graph_node_by_name(model.graph, node) args = output_node.args[0] if isinstance(args, torch.fx.Node): + # Case of non tuple output. args = value_remap[args] else: + # Case of tuple output. args = [value_remap[n] for n in args] # Unpack target output args in case - # only one arg is presented. + # the only one arg is presented. if len(args) == 1: args = args[0] nodes_with_output[idx] = args