Merge branch 'main' into datasets/eurosat

pytorch · Aug 6, 2024 · e135137 · e135137
2 parents 2f33f46 + c8c496d
commit e135137
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/torchvision/csrc/ops/autocast/nms_kernel.cpp b/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
diff --git a/torchvision/csrc/ops/autocast/roi_align_kernel.cpp b/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