Fix CPU segfault in roi_align when ROI contains NaN coordinates

[nileshpatil6]

Fixes #9273

roi_align on CPU can segfault when a ROI has NaN coordinates. I tracked it down to the bounds check in pre_calc_for_bilinear_interpolate (torchvision/csrc/ops/cpu/roi_align_common.h) and the equivalent check in bilinear_interpolate_gradient (torchvision/csrc/ops/cpu/roi_align_kernel.cpp, used by the backward pass).

Both do something like:

if (y < -1.0 || y > height || x < -1.0 || x > width) {
  // treat as empty ROI
}

The problem is that any comparison against NaN is false, so a NaN coordinate slides right through this check instead of being caught as out of bounds. Right after that, the code does (int)y to compute the pixel index, and casting NaN to int is undefined behavior in C++. In practice it produces some garbage integer, which then gets used as pos1/pos2/pos3/pos4 to read straight out of the input feature map buffer, which is exactly what the ASAN report in the issue shows.

I reproduced the crash locally with the repro script from the issue against the currently released torchvision (0.21.0) and it segfaults reliably.

Fix is small: add explicit isnan checks next to the existing range checks in both places, so a NaN ROI gets treated the same way as any other malformed/out-of-bounds ROI (zero weights, index -1, no memory read), instead of silently producing garbage.

Added a regression test (test_nan_roi_coordinates_do_not_crash in test/test_ops.py) that runs roi_align with a NaN ROI on CPU for both aligned values and both sampling_ratio modes.

One thing I want to be upfront about: I don't have a C++ toolchain available in the environment I used to write this patch, so I could not actually compile torchvision from source to run the new test against my fix. What I did verify:

• The crash is real: I reproduced the exact segfault from the issue against the released 0.21.0 wheel.
• The fix logic itself: I ported the exact before/after C++ conditions to Python and confirmed the old condition returns False for NaN inputs (bug) and the new one returns True (fixed), matching the intended behavior change.

If a maintainer or CI can confirm the actual build + test pass, that'd be great, happy to iterate if anything doesn't compile as expected (e.g. if isnan needs a different form for the Half/BFloat16 specializations dispatched via AT_DISPATCH_FLOATING_TYPES_AND_HALF).

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9541

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.