A bug for demo to infer one image

[zifuwan]

Traceback (most recent call last):
File "try_mask_auto_gen.py", line 5, in
masks = mask_generator.generate(input_image)
File "/home/zifu/anaconda3/envs/sam/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/home/zifu/Semantic-SAM/tasks/automatic_mask_generator.py", line 182, in generate
mask_data = self._generate_masks(image)
File "/home/zifu/Semantic-SAM/tasks/automatic_mask_generator.py", line 225, in _generate_masks
crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)
File "/home/zifu/Semantic-SAM/tasks/automatic_mask_generator.py", line 267, in _process_crop
data.cat(batch_data)
File "/home/zifu/Semantic-SAM/utils/sam_utils/amg.py", line 65, in cat
self._stats[k] = torch.cat([self._stats[k], v], dim=0)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument tensors in method wrapper_cat)

To solve this, go to Semantic-SAM/utils/sam_utils/amg.py and modify like:
def cat(self, new_stats: "MaskData") -> None:
for k, v in new_stats.items():
if k not in self._stats or self._stats[k] is None:
self._stats[k] = deepcopy(v)
elif isinstance(v, torch.Tensor):
self._stats[k] = torch.cat([self._stats[k].cuda(), v.cuda()], dim=0)
elif isinstance(v, np.ndarray):
self._stats[k] = np.concatenate([self._stats[k], v], axis=0)
elif isinstance(v, list):
self._stats[k] = self._stats[k] + deepcopy(v)
else:
raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.")

[linuaries]

I got the same error, and your modification solve this problem. Thanks.