pytorch
213e10dc - fix bug in trace model when out-operator has more than one output (#101563)

Commit
1 year ago
fix bug in trace model when out-operator has more than one output (#101563) Fixes #https://github.com/pytorch/pytorch/issues/101960 when I trace a func to run out-operator has more than one output, I got the error. This is because the situation when the output of the out operator is greater than 1 is not handled. ``` def test_trace_out_operator_with_two_output(): example_input = torch.rand(2, 8) out_1, out_2 = torch.cummax(example_input, 1) def run_cummax(example_input, out_1, out_2): output_1, output_2 = torch.cummax(example_input, 1, out=(out_1, out_2)) return output_1, output_2 trace_model = torch.jit.trace(run_cummax, (example_input, out_1, out_2)) and the error info: raise TracingCheckError( torch.jit._trace.TracingCheckError: Tracing failed sanity checks! encountered an exception while running the trace with test inputs ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/101563 Approved by: https://github.com/jgong5, https://github.com/EikanWang, https://github.com/davidberard98
Committer
Parents
Loading