Bug in rapid.py

[Alex18947]

Hi, while training with custom fisheye dataset on a pretrained COCO checkpoint, I encounter some "index out of bounds" problems in rapid.py: here

I am running with CUDA_LAUNCH_BLOCKING=1 to get the exact line. The error is:

./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.

Adding some prints before the error I can see that penalty_mask is of size: torch.Size([4, 3, 136, 136])

Other values:

b: 2 best_n: tensor([0, 0, 0, 0, 2]) truth_j: tensor([141, 99, 47, 53, 97]) truth_i: tensor([ 81, 109, 86, 96, 108])

I don't seem able to find the root cause as for now, seems like it happens for some images only. Any ideas?