torchvision.roi_align does not support TPU

[CheonJiEun]

Hello.
We are using TPU in GCP.

We are currently modifying the code to allow the TPU to return to Detectron2.
However, there is an error that roi_align in Torchvision is not supported by TPU.
Please check the bottom. Can you solve it for me?

File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torchvision/ops/roi_align.py", line 51, in roi_align return torch.ops.torchvision.roi_align(input, rois, spatial_scale, output_size[0], output_size[1], sampling_ratio, aligned) RuntimeError: Could not run 'torchvision::roi_align' with arguments from the 'XLA' backend. 'torchvision::roi_align' is only available for these backends: [CPU, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, Tracer, Autocast, Batched, VmapMode].

[vfdev-5]

@CheonJiEun currently torchvision do not provide any support for TPU. Certain ops (like transforms) could work out-of-the-box, but others not.
cc @fmassa for more details.