Torch onnx (#48980)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/45215
This is a follow up PR of https://github.com/pytorch/pytorch/issues/45258 and https://github.com/pytorch/pytorch/issues/48782
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48980
Reviewed By: zhangguanheng66
Differential Revision: D25399823
Pulled By: ezyang
fbshipit-source-id: 798055f4abbbffecdfab0325884193c81addecec
