[fx2trt] make gpu trace better (#65168)
authorShiyan Deng <dsy842974287@fb.com>
Fri, 17 Sep 2021 16:23:05 +0000 (09:23 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Fri, 17 Sep 2021 16:32:37 +0000 (09:32 -0700)
commit6aa04b68436ea4684e530baa9bd91abe2d629306
tree761d3194aa55ee4cc68655754587a052219b24a1
parenta8d7b885c52e9acd25775bbd50de238da5a5624d
[fx2trt] make gpu trace better (#65168)

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

Add record_function to TRTModule and EngineHolder so each parts would appear on gpu trace.

Test Plan: CI

Reviewed By: wushirong

Differential Revision: D30997968

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