Torchvision Object detection TPU Support

❓ Torchvision object detection models with TPU.

My doubt lies somewhere between feature request and question. hence posting here.

PyTorch supports TPU through torch_xla. It makes it possible to train models over TPU.
I guess most torchvision classification models can be used with transfer learning/training over TPU.

For torchvision object detection models, do they support TPU?
Some operations such as NMS, rpn, roi_align do not support TPU and hence I get an error as follows.

I was trying Faster R-CNN resnet50 fpn model for object detection.

  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torchvision/models/detection/generalized_rcnn.py", line 70, in forward
    proposals, proposal_losses = self.rpn(images, features, targets)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torchvision/models/detection/rpn.py", line 493, in forward
    boxes, scores = self.filter_proposals(proposals, objectness, images.image_sizes, num_anchors_per_level)
  File "/usr/local/lib/python3.6/dist-packages/torchvision/models/detection/rpn.py", line 416, in filter_proposals
    keep = box_ops.batched_nms(boxes, scores, lvl, self.nms_thresh)
RuntimeError: Cannot access data pointer of Tensor that doesn't have storage

My doubts/concerns/feature request.

1. Do torchvision object detection models support TPU training?
2. Any Plans for TPU support in future releases for these models?
3. Are these ops only CUDA native and GPU/CPU specific? Is there a work-around to train object detection / segmentation models with TPU?

Start by reviewing torchvision's Faster R-CNN path, especially detection/rpn.py and the box_ops.batched_nms call shown in the traceback. Check how torch_xla handles NMS, RPN, and roi_align, and whether the reported storage-pointer error has an existing workaround. Done would require a documented or implemented path for TPU training, including support for the affected operations.