Support .to(device) or Device Aware Handling for Segmentation Labels in EOMTImageProcessor #42205

src/transformers/models/eomt/modeling_eomt.py

@@ -1104,6 +1104,14 @@ def forward(
             list of tuples indicating the image index and start and end positions of patches for semantic segmentation.
         """
 
+        if mask_labels is not None:
+            target_device = pixel_values.device
+            mask_labels = [mask.to(target_device) for mask in mask_labels]
+
+        if class_labels is not None:
+            target_device = pixel_values.device
+            class_labels = [label.to(target_device) for label in class_labels]
+
         masks_queries_logits_per_layer, class_queries_logits_per_layer = (), ()
         attention_mask = None
 

src/transformers/models/eomt/modular_eomt.py

@@ -513,6 +513,14 @@ def forward(
             list of tuples indicating the image index and start and end positions of patches for semantic segmentation.
         """
 
+        if mask_labels is not None:
+            target_device = pixel_values.device
+            mask_labels = [mask.to(target_device) for mask in mask_labels]
+
+        if class_labels is not None:
+            target_device = pixel_values.device
+            class_labels = [label.to(target_device) for label in class_labels]
+
         masks_queries_logits_per_layer, class_queries_logits_per_layer = (), ()
         attention_mask = None
 

src/transformers/models/mask2former/modeling_mask2former.py

@@ -2415,6 +2415,15 @@ def forward(
         torch.Size([338, 676])
         ```
         """
+
+        if mask_labels is not None:
+            target_device = pixel_values.device
+            mask_labels = [mask.to(target_device) for mask in mask_labels]
+
+        if class_labels is not None:
+            target_device = pixel_values.device
+            class_labels = [label.to(target_device) for label in class_labels]
+
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states

src/transformers/models/maskformer/modeling_maskformer.py

@@ -1739,6 +1739,13 @@ def forward(
         [480, 640]
         ```
         """
+        if mask_labels is not None:
+            target_device = pixel_values.device
+            mask_labels = [mask.to(target_device) for mask in mask_labels]
+
+        if class_labels is not None:
+            target_device = pixel_values.device
+            class_labels = [label.to(target_device) for label in class_labels]
 
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (

src/transformers/models/oneformer/modeling_oneformer.py

@@ -3141,6 +3141,13 @@ def forward(
         '👉 Panoptic Predictions Shape: [512, 683]'
         ```
         """
+        if mask_labels is not None:
+            target_device = pixel_values.device
+            mask_labels = [mask.to(target_device) for mask in mask_labels]
+
+        if class_labels is not None:
+            target_device = pixel_values.device
+            class_labels = [label.to(target_device) for label in class_labels]
 
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (