facebookresearch / segment-anything-2

The repository provides code for running inference with the Meta Segment Anything Model 2 (SAM 2), links for downloading the trained model checkpoints, and example notebooks that show how to use the model.
Apache License 2.0
10.76k stars 874 forks source link

Adding gpu box tensors and no points in add_new_points_or_box does not work #281

Open s-fleuren opened 2 weeks ago

s-fleuren commented 2 weeks ago

When passing a tensor that resides on the gpu to the box argument of add_new_points_or_box and not adding any points, the function raises an error:

File "/home2/steven/Software/segment-anything-2/sam2/sam2_video_predictor.py", line 226, in add_new_points_or_box points = torch.cat([box_coords, points], dim=1) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument tensors in method wrapper_CUDA_cat)

This happens because the points tensor is initialized on the cpu when the points argument is none:

if points is None:
            points = torch.zeros(0, 2, dtype=torch.float32)