Last fix

osrt/sam/prompt_encoder.py

@@ -81,9 +81,11 @@ class PromptEncoder(nn.Module):
         if pad:
             padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
             padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
-            points = torch.cat([points, padding_point], dim=1)
-            labels = torch.cat([labels, padding_label], dim=1)
+            points = torch.cat([points, padding_point], dim=1).to(points.device)
+            labels = torch.cat([labels, padding_label], dim=1).to(labels.device)
         point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size)
+        self.not_a_point_embed.to(points.device)
+        self.point_embeddings.to(points.device)
         point_embedding[labels == -1] = 0.0
         point_embedding[labels == -1] += self.not_a_point_embed.weight
         point_embedding[labels == 0] += self.point_embeddings[0].weight