RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu) using Crop N layers > 1

[bthetford]

Full traceback:

Exception in thread Thread-1 (run):
Traceback (most recent call last):
  File "C:\Users\M.Modeler\AppData\Local\miniconda3\envs\geoSieve\Lib\threading.py", line 1038, in _bootstrap_inner
    self.run()
  File "C:\Users\M.Modeler\AppData\Local\miniconda3\envs\geoSieve\Lib\threading.py", line 975, in run
    self._target(*self._args, **self._kwargs)
  File "C:\Users\M.Modeler\geoSieve\guiUtil.py", line 1037, in run
    self.fun(self.settings)
  File "C:\Users\M.Modeler\geoSieve\backend.py", line 255, in geoSieve
    sam_dat = util.run_sam(img,
              ^^^^^^^^^^^^^^^^^
  File "C:\Users\M.Modeler\geoSieve\util.py", line 262, in run_sam
    sam_outputs = mask_generator.generate(image)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\M.Modeler\AppData\Local\miniconda3\envs\geoSieve\Lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\M.Modeler\AppData\Local\miniconda3\envs\geoSieve\Lib\site-packages\segment_anything\automatic_mask_generator.py", line 163, in generate
    mask_data = self._generate_masks(image)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\M.Modeler\AppData\Local\miniconda3\envs\geoSieve\Lib\site-packages\segment_anything\automatic_mask_generator.py", line 214, in _generate_masks
    keep_by_nms = batched_nms(
                  ^^^^^^^^^^^^
  File "C:\Users\M.Modeler\AppData\Local\miniconda3\envs\geoSieve\Lib\site-packages\torchvision\ops\boxes.py", line 73, in batched_nms
    return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\M.Modeler\AppData\Local\miniconda3\envs\geoSieve\Lib\site-packages\torch\jit\_trace.py", line 1220, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "C:\Users\M.Modeler\AppData\Local\miniconda3\envs\geoSieve\Lib\site-packages\torchvision\ops\boxes.py", line 110, in _batched_nms_vanilla
    keep_mask[curr_indices[curr_keep_indices]] = True
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Was able to fix this by adding device = "cuda" to the following call of torch.zeros() on line 217 of automatic_mask_generator.py:

        if len(crop_boxes) > 1:
            # Prefer masks from smaller crops
            scores = 1 / box_area(data["crop_boxes"])
            scores = scores.to(data["boxes"].device)
            keep_by_nms = batched_nms(
                data["boxes"].float(),
                scores,
                torch.zeros(len(data["boxes"]), device = "cuda"),  # categories
                iou_threshold=self.crop_nms_thresh,
            )
            data.filter(keep_by_nms)

I noticed I'm using an old package version right now, and that the updated script now uses torch.zeros_like() here, but this might still apply. Of course my workaround doesn't allow for any other devices other than cuda.

[bthetford]

This was fixed by the updated code which uses torch.zeroes_like here instead of torch.zeroes