Tensor.nonzero fails on GPU for tensors containing more than INT_MAX elements

[MartinKlefas]

I've been gradually increasing the size of the input images I'm using with Segment Anything, and it seems as though once you have an image where the number of detections x height * width > 2^31 you run into an unresolved torch bug:
pytorch/pytorch#51871

I've managed to create this using an image that's 4096 x 4096 pixels and with setting that detect > 128 objects.
I get the error message:

Traceback (most recent call last):
File "/venv/lib/python3.10/site-packages/runpod/serverless/modules/rp_job.py", line 113, in run_job
result = handler(job)
File "/make_crypts.py", line 97, in runpod_mask_save
thisMaskSet = mask_generator.generate(patch_img)
File "/venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/venv/lib/python3.10/site-packages/segment_anything/automatic_mask_generator.py", line 163, in generate
mask_data = self._generate_masks(image)
File "/venv/lib/python3.10/site-packages/segment_anything/automatic_mask_generator.py", line 206, in _generate_masks
crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)
File "/venv/lib/python3.10/site-packages/segment_anything/automatic_mask_generator.py", line 245, in _process_crop
batch_data = self._process_batch(points, cropped_im_size, crop_box, orig_size)
File "/venv/lib/python3.10/site-packages/segment_anything/automatic_mask_generator.py", line 318, in _process_batch
data["rles"] = mask_to_rle_pytorch(data["masks"])
File "/venv/lib/python3.10/site-packages/segment_anything/utils/amg.py", line 118, 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

One resolution would be to anticipate this image and process only a number of masks at a time that would avoid hitting INT_MAX

Another would be to use 64 bit integers - though I don't know if that's possible in Torch.

Please let me know if you need any further information or example code.

[nikky4D]

I've come across this error. Is there any development or a resolution on your end?

[MartinKlefas]

I've come across this error. Is there any development or a resolution on your end?

The error is mostly caused in an external library, but I've created a workaround - until it's merged you can use my git fork, though it may be better security practice to fork your own to make similar changes