Fix oneformer instance segmentation RuntimeError (huggingface#27725)

src/transformers/models/oneformer/image_processing_oneformer.py

@@ -1167,6 +1167,7 @@ def post_process_instance_segmentation(
  class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1]
  masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width]
 
+ device = masks_queries_logits.device
  batch_size = class_queries_logits.shape[0]
  num_queries = class_queries_logits.shape[1]
  num_classes = class_queries_logits.shape[-1] - 1
@@ -1177,7 +1178,7 @@ def post_process_instance_segmentation(
  for i in range(batch_size):
  # [Q, K]
  scores = torch.nn.functional.softmax(class_queries_logits[i], dim=-1)[:, :-1]
- labels = torch.arange(num_classes).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1)
+ labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1)
 
  # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
  scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False)