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

Signed-off-by: Feng Yuan <[email protected]>
Co-authored-by: Nicolas Hug <[email protected]>

Author: fengyuan14

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