Fix TRTModule not adding outputs in order (#64418)
authorKefei Lu <kefeilu@fb.com>
Thu, 2 Sep 2021 08:17:56 +0000 (01:17 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Thu, 2 Sep 2021 08:36:23 +0000 (01:36 -0700)
commit6db8f7a70920f91418078fe09477eed0b0adefdb
tree04caa82f11d87829a691ea862647a9b879b62c66
parent76e187aa08556ce90e84b17e836784ffbb6905e0
Fix TRTModule not adding outputs in order (#64418)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64418

In T99368564, we found that when running TRT lowered module, the output tensors are out-of-order, as compared to the output from the original, non-lowered module. It turns out that in `TRTModule.forward()`, we cannot rely on `ICudaEngine` bindings natural order indices to create the output tensors, but rather, we should explicitly construct the output tensor from the bindings' names, in an ordered that we supply.

Test Plan:
* Arc lint
* Run CI/sandcastle tests
* Run GPU lowering using commands and code changes in D30171741 and ensure we don't observe out-of-order outputs

Reviewed By: yinghai

Differential Revision: D30693545

fbshipit-source-id: 32a894ceeb148fcf4e8d279be3835c7d1f1aa2ba
torch/fx/experimental/fx2trt/fx2trt.py