trevor-m opened a new pull request #5549: URL: https://github.com/apache/incubator-tvm/pull/5549
This bug occurs when an op has multiple outputs in pytorch, but only has a single output in relay. For example, this bug occurs with `aten::max_pool2d_with_indices` which has two outputs but relay's max_pool only has one: `Tensor, %34 : Tensor = aten::max_pool2d_with_indices(%input_4.1, %29, %30, %31, %32, %3) ` [This condition](https://github.com/apache/incubator-tvm/blob/master/python/tvm/relay/frontend/pytorch.py#L1923) is hit: ``` if node.outputsSize() > 1: node_name = "_".join(_get_output_names(node)) ``` which gives this node the mangled name of `input_5.1_34`, instead of `input_5.1`. Since the TVM output is not a tuple, this mangled name is used as the key in the outputs dict. https://github.com/apache/incubator-tvm/blob/master/python/tvm/relay/frontend/pytorch.py#L2238-L2244 Later on, when the output of the node is needed, `input_5.1` cannot be found in the dict, causing imprting the model to fail. ``` File ".../tvm/relay/frontend/pytorch.py", line 1781, in <listcomp> return [outputs[name] for name in _get_input_names(op_node)] KeyError: 'input_5.1' ``` This PR fixes the bug by using the first output name to store the node output in the outputs dict. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org
