System Info

• transformers version: 4.53.1
• Platform: Linux-6.8.0-86-generic-x86_64-with-glibc2.35
• Python version: 3.10.12
• Huggingface_hub version: 0.33.2
• Safetensors version: 0.5.3
• Accelerate version: 1.8.1
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (accelerator?): 2.7.1+cu126 (CUDA)
• Tensorflow version (GPU?): 2.19.0 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?:
• Using GPU in script?:
• GPU type: NVIDIA GeForce RTX 4050 Laptop GPU

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Description

When processing keypoint pairs with different numbers of valid keypoints, the current post-processing logic can attempt to index out of bounds on the second keypoint tensor (keypoints1). This happens because invalid match indices are not checked against mask1 before indexing.

Root Cause

Masks are applied to the keypoints, preserving only valid correspondences:

keypoints0 = keypoints_pair[0][mask0]
keypoints1 = keypoints_pair[1][mask1]

This generates potentially inconsistent dimensions, e.g. len(keypoints0) = 88 but len(keypoints1) = 87.

At this point the code filters the matches considering only the first mask:

matches0 = matches[mask0]
scores0 = scores[mask0]

The method then filters out matches with low scores, but not the matches that are still pointing to out of bounds indices in the previously shortened keypoints1:

valid_matches = torch.logical_and(scores0 > threshold, matches0 > -1)

Example

Bug encountered with the following pair of images:

Consider the case where:

mask1 = tensor([T, T, T, T, T, T, ...., T, F])  # shape: [88]
keypoints1 = keypoints_pair[1][mask1]            # shape: [87, 2]
matches0 = tensor([87, -1, -1, 3, 10, -1, 5])   # shape: [88]
                   ↑ Index 87 doesn't exist in keypoints1 (only indices 0-86)

This causes an unhandled out of bounds exception here:

matched_keypoints1 = keypoints1[matches0[valid_matches]]

CUDA error: device-side assert triggered
../aten/src/ATen/native/cuda/IndexKernel.cu:93: operator(): block: [0,0,0], thread: [0,0,0]
Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.

Expected behavior

Proposed Fix

mask0 = mask_pair[0] > 0
mask1 = mask_pair[1] > 0
keypoints0 = keypoints_pair[0][mask0]
keypoints1 = keypoints_pair[1][mask1]
matches0 = matches[mask0]
scores0 = scores[mask0]

# FIX: Verify that matches point to valid (unmasked) keypoints in mask1
valid_in_mask1 = torch.zeros_like(matches0, dtype=torch.bool)
for i, match_idx in enumerate(matches0):
    if match_idx > -1 and match_idx < len(mask1) **and mask1[match_idx]:**
        valid_in_mask1[i] = True

# Filter out matches with low scores OR pointing to masked keypoints
valid_matches = torch.logical_and(scores0 > threshold, valid_in_mask1)

matched_keypoints0 = keypoints0[valid_matches]
matched_keypoints1 = keypoints1[matches0[valid_matches]]
matching_scores = scores0[valid_matches]