Closed
Description
🚀 Feature
Currently only NMS and RoIAlign support autocast (introduced in PyTorch 1.6.0). We should extend support for autocast to all C++ ops in torchvision.
Motivation
Mixed precision training is PyTorch is obtained via support for autocast (automatically casting between fp16 -> fp32 depending on the operation).
The work should follow what was done for the RoIAlign function in #2384 (as that PR also converted NMS to use the dispatcher).
Adding support for autocast means that the op should be registered via the dispatcher, so this is blocked by #2796
Pitch
Out-of-the-box mixed precision support for all operators implemented in torchvision:
- nms
- roi_align
- deform_conv2d
- roi_pool
- ps_roi_align
- ps_roi_pool