(feat): Moving labels to same device as logits for Deit (#22679)

src/transformers/models/deit/modeling_deit.py

@@ -764,6 +764,7 @@ class DeiTForImageClassification(DeiTPreTrainedModel):
 
         loss = None
         if labels is not None:
+            labels = labels.to(logits.device)
             if self.config.problem_type is None:
                 if self.num_labels == 1:
                     self.config.problem_type = "regression"