Setting multiple points prompts results in an error

[jcorsetti]

Hello, thanks for open-sourcing the code and for the great work. I am trying to use RobustSam to produce N segmentation masks given N point prompts, in a single image with a single forward pass. From what I understood this is possible from the comments to the forward function here:

https://github.com/robustsam/RobustSAM/blob/6f0e2637e58f812985cdf9f403a5d7515a66fe38/robust_segment_anything/modeling/sam.py#L74C1-L74C71

However, when I try to do so an error occurs in the mask decoder. I tried a couple of different options, and it seems to work only with a single point prompt. In my case (5 points), it throws the following error. Is this the intended behaviour or am I missing something? In the following screenshot you can see the shapes of my input tensors and the error trace.