Fix arg type annotations in jit.trace and onnx.export (#41093)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/40350
Pull Request resolved: https://github.com/pytorch/pytorch/pull/41093
Differential Revision: D22477950
Pulled By: malfet
fbshipit-source-id: f1141c129b6d9efb373d22291b441df86c529ddd
