Merge branch 'main' into adabeyta/ps-roi-pool-cuda-stable

Author: NicolasHug

CMakeLists.txt

@@ -116,6 +116,7 @@ set(TORCHVISION_STABLE_SOURCES
   ${TVCPP}/ops/cuda/ps_roi_pool_kernel.cu
   ${TVCPP}/ops/ps_roi_align.cpp
   ${TVCPP}/ops/cpu/ps_roi_align_kernel.cpp
+  ${TVCPP}/ops/cuda/ps_roi_align_kernel.cu
   ${TVCPP}/ops/deform_conv2d.cpp
   ${TVCPP}/ops/cpu/deform_conv2d_kernel.cpp
   ${TVCPP}/io/image/cuda/decode_jpegs_cuda.cpp

setup.py

@@ -188,6 +188,7 @@ def get_macros_and_flags():
 STABLE_SOURCES.add(CSRS_DIR / ("ops/hip/roi_pool_kernel.hip" if IS_ROCM else "ops/cuda/roi_pool_kernel.cu"))
 STABLE_SOURCES.add(CSRS_DIR / ("ops/hip/roi_align_kernel.hip" if IS_ROCM else "ops/cuda/roi_align_kernel.cu"))
 STABLE_SOURCES.add(CSRS_DIR / ("ops/hip/ps_roi_pool_kernel.hip" if IS_ROCM else "ops/cuda/ps_roi_pool_kernel.cu"))
+STABLE_SOURCES.add(CSRS_DIR / ("ops/hip/ps_roi_align_kernel.hip" if IS_ROCM else "ops/cuda/ps_roi_align_kernel.cu"))
 
 
 def _not_stable(paths):

torchvision/_autograd_registrations.py

@@ -87,7 +87,7 @@ def _ps_roi_align_setup_context(ctx, inputs, output):
 
 
 def _ps_roi_align_backward(ctx, grad_output, _grad_channel_mapping):
-    if grad_output.device.type == "mps":
+    if grad_output.device.type in ("cuda", "mps"):
         torch._prims_common.alert_not_deterministic("ps_roi_align_backward_kernel")
     rois, channel_mapping = ctx.saved_tensors
     batch_size, channels, height, width = ctx.input_shape