error with shape size

[wangmiaowei]

I found the shape size can also cause bug, when my input is with the size image: (3006, 4013, 3).

I get error like:
RuntimeError: nonzero is not supported for tensors with more than INT_MAX elements, file a support request:
the full list:

warnings.warn(
Traceback (most recent call last):
File "/home/wmw/project/codespace/segment-anything-2-main/test.py", line 44, in
masks = mask_generator.generate(image)
File "/home/wmw/miniconda3/envs/gspl/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/home/wmw/project/codespace/segment-anything-2-main/sam2/automatic_mask_generator.py", line 178, in generate
mask_data = self._generate_masks(image)
File "/home/wmw/project/codespace/segment-anything-2-main/sam2/automatic_mask_generator.py", line 215, in _generate_masks
crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)
File "/home/wmw/project/codespace/segment-anything-2-main/sam2/automatic_mask_generator.py", line 253, in _process_crop
batch_data = self._process_batch(
File "/home/wmw/project/codespace/segment-anything-2-main/sam2/automatic_mask_generator.py", line 361, in _process_batch
data["rles"] = mask_to_rle_pytorch(data["masks"])
File "/home/wmw/project/codespace/segment-anything-2-main/sam2/utils/amg.py", line 120, in mask_to_rle_pytorch
change_indices = diff.nonzero()
RuntimeError: nonzero is not supported for tensors with more than INT_MAX elements, file a support request

[bhack]

It seems really related to pytorch/pytorch#125850

[wangmiaowei]

yeas, so no solution currently?

[heyoeyo]

The same problem occurred on SAM v1 and someone posted a change to the code that fixed it:
facebookresearch/segment-anything#569

I assume something similar would work with v2. Alternatively, I believe running on CPU also avoids the problem.

[bhack]

Yes with that workaround it will work but it is going to be slower as you need to process it in chunks.
It is an upstream CUDA/CUB limit as pytorch nonzero is implemented on Nvidia GPU using CUB primitives.
It seems that Nvidia it is still lagging to fix it as you can see from the upstream ticket:
NVIDIA/cccl#1422