roi_align and roi_pool cannot be scripted due to type annotation issue

[neginraoof]

🚀 Feature

roi_align anf roi_pool cannot be scripted due to type annotation issue

vision/torchvision/ops/roi_pool.py

Line 13 in 7bc5bd8

boxes: Tensor,

vision/torchvision/ops/roi_align.py

Line 13 in 7bc5bd8

boxes: Tensor,

The input "boxes" can be either a tensor or list, however the type annotation does not cover the list case.
Once JIT support Union types the annotation should be updated to Union of Tensor and List.

[datumbox]

JIT does not support unions at the moment. That's why it's only defined as Tensor. See: pytorch/pytorch#41412

Why don't you pass a tensor instead?

[neginraoof]

We can modify the model code. But this issue has been escalated to the ONNX team a couple times by customers who are relying on the documentation. So we'd like to track this work and make sure this is fixed once Union is supported.

[datumbox]

@neginraoof Sounds good. There are numerous places where things will change once Union is supported and I think this use-case will be one of them. Perhaps if you rephrase the ticket as a feature request / enhancement dependent on JIT Union support, I can tag it differently and leave it open. The challenge right now is it states it's a bug which means if someone tries to "fix" it by putting Union it will definitely break JIT completely.

[fmassa]

I think what we can do for now is to improve the documentation to mention that torchscript only supports Tensor, while eager mode supports both.

It looks like Union would be soon supported in torchscript though, so we might instead just wait a bit more before we modify the type annotations

[pmeier]

pytorch/pytorch#53180 for reference.

[neginraoof]

Thanks @fmassa, we're okay with waiting for Union to go in.

@datumbox I updated the description.