[fbsync] Register nms and roi_align Autocast policy for PyTorch Intel GPU backend (#8541)

Summary: Signed-off-by: Feng Yuan <[email protected]>

Differential Revision: D60903715

fbshipit-source-id: c0919c6b0473432b7ef7a30599c3dc143e3da1cf

Co-authored-by: Nicolas Hug <[email protected]>

Author: NicolasHug

torchvision/csrc/ops/autocast/nms_kernel.cpp

@@ -38,5 +38,12 @@ TORCH_LIBRARY_IMPL(torchvision, AutocastCPU, m) {
           (nms_autocast<c10::DispatchKey::AutocastCPU, c10::DeviceType::CPU>)));
 }
 
+TORCH_LIBRARY_IMPL(torchvision, AutocastXPU, m) {
+  m.impl(
+      TORCH_SELECTIVE_NAME("torchvision::nms"),
+      TORCH_FN(
+          (nms_autocast<c10::DispatchKey::AutocastXPU, c10::DeviceType::XPU>)));
+}
+
 } // namespace ops
 } // namespace vision

torchvision/csrc/ops/autocast/roi_align_kernel.cpp

@@ -48,5 +48,13 @@ TORCH_LIBRARY_IMPL(torchvision, AutocastCPU, m) {
                 c10::DeviceType::CPU>)));
 }
 
+TORCH_LIBRARY_IMPL(torchvision, AutocastXPU, m) {
+  m.impl(
+      TORCH_SELECTIVE_NAME("torchvision::roi_align"),
+      TORCH_FN((roi_align_autocast<
+                c10::DispatchKey::AutocastXPU,
+                c10::DeviceType::XPU>)));
+}
+
 } // namespace ops
 } // namespace vision