Feature Proposal: -
Torchvision Object detection models should support TPU.
TPUs will greatly lower training time for models like Faster RCNN and coming soon Retina Net. Torch natively supports TPU through xla library. It is possible to train the classification models on TPU as well. But when we try object detection models on TPU it gives the following error
 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

This error is because some features such rpn.py , ops.batched_nms is not TPU compatible.
Will we get any support for TPU soon for these models ?
I created an issue 9 on GitHub as well but I got no response. So I thought maybe I will add here to discuss.