"IndexError: _Map_base::at" when calling CropAndResize

[nathanpainchaud]

🐛 Bug

I updated the conda environment I use for my project (based on PyTorch Lightning and PyTorch), and I'm now getting an error when making the following call (where x is a Tensor):

from roi_align import CropAndResize
CropAndResize(height, width)(x, roi_bbox_hat, torch.arange(x.shape[0], dtype=torch.int32, device=x.device))

The error I'm getting is the following:

Traceback (most recent call last):
  [...]
  File "my_code.py", line XX, in forward
    cropped_x = CropAndResize(height, width)(x, roi_bbox_hat, torch.arange(x.shape[0], dtype=torch.int32, device=x.device))
  File "$CONDA_ENV_PATH/lib/python3.8/site-packages/torch/nn/modules/module.py", line 720, in _call_impl
    result = self._slow_forward(*input, **kwargs)
  File "$CONDA_ENV_PATH/lib/python3.8/site-packages/torch/nn/modules/module.py", line 704, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "$CONDA_ENV_PATH/lib/python3.8/site-packages/roi_align/crop_and_resize.py", line 70, in forward
    return CropAndResizeFunction.apply(image, boxes, box_ind, self.crop_height, self.crop_width, self.extrapolation_value)
  File "$CONDA_ENV_PATH/lib/python3.8/site-packages/roi_align/crop_and_resize.py", line 23, in forward
    crop_and_resize_gpu.forward(
IndexError: _Map_base::at
python-BaseException

Prior to the update, I was using PyTorch Lightning 0.8.5 with PyTorch 1.5. I assume the cause of the error is related to the update from PyTorch 1.5 to 1.6 (since it's a major update), but I can't pinpoint the exact cause. Does anyone have clues about the cause of the error and/or suggestions on how to fix it?

Environment

Below is my configuration following the update (I didn't change any of the hardware).

• CUDA:
  - GPU: TITAN Xp
  - available: True
  - version: 10.2
• Packages:
  - numpy: 1.19.1
  - pyTorch_debug: False
  - pyTorch_version: 1.6.0
  - pytorch-lightning: 0.9.0
  - tensorboard: 2.2.0
  - tqdm: 4.48.2
• System:
  - OS: Linux
  - architecture: 64bit, ELF
  - processor: x86_64
  - python: 3.8.5
  - version: #56~18.04.1-Ubuntu SMP