[JIT] script if tracing fix #40468
[JIT] script if tracing fix #40468
Conversation
There was a problem hiding this comment.
@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
facebook-github-bot
added
the
oncall: jit
💊 CI failures summary and remediationsAs of commit fcfd2e6 (more details on the Dr. CI page): ✅ None of the CI failures appear to be your fault 💚
🚧 1 fixed upstream failure:These were probably caused by upstream breakages that were already fixed.
Please rebase on the
|
| return fn(*args) | ||
| return fn(*args, **kwargs) | ||
|
|
||
| compiled_fn = script(wrapper.__original_fn) |
There was a problem hiding this comment.
caching already occurs when we compile a fn, so this was unnecessary
There was a problem hiding this comment.
@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
There was a problem hiding this comment.
@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
Summary: Currently, torchvision annotates `batched_nms` with `torch.jit.script` so the function gets compiled when it is traced and ONNX will work. Unfortunately, this means we are eagerly compiling batched_nms, which fails if torchvision isn't built with `torchvision.ops.nms`. As a result, torchvision doesn't work on torch hub right now. `_script_if_tracing` could solve our problem here, but right now it does not correctly interact with recursive compilation. This PR fixes that bug. Pull Request resolved: pytorch#40468 Reviewed By: jamesr66a Differential Revision: D22195771 Pulled By: eellison fbshipit-source-id: 83022ca0bab6d389a48a478aec03052c9282d2b7
Summary: Currently, torchvision annotates `batched_nms` with `torch.jit.script` so the function gets compiled when it is traced and ONNX will work. Unfortunately, this means we are eagerly compiling batched_nms, which fails if torchvision isn't built with `torchvision.ops.nms`. As a result, torchvision doesn't work on torch hub right now. `_script_if_tracing` could solve our problem here, but right now it does not correctly interact with recursive compilation. This PR fixes that bug. Pull Request resolved: #40468 Reviewed By: jamesr66a Differential Revision: D22195771 Pulled By: eellison fbshipit-source-id: 83022ca0bab6d389a48a478aec03052c9282d2b7 Co-authored-by: Elias Ellison <eellison@fb.com>
Currently, torchvision annotates
batched_nmswithtorch.jit.scriptso the function gets compiled when it is traced and ONNX will work. Unfortunately, this means we are eagerly compiling batched_nms, which fails if torchvision isn't built withtorchvision.ops.nms. As a result, torchvision doesn't work on torch hub right now._script_if_tracingcould solve our problem here, but right now it does not correctly interact with recursive compilation. This PR fixes that bug.