moved labels to the same device as logits for LILT model (huggingface…

…#22898)
novice03 · Jun 23, 2023 · e6c3db1 · e6c3db1
1 parent c1ebdae
commit e6c3db1
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/transformers/models/lilt/modeling_lilt.py b/src/transformers/models/lilt/modeling_lilt.py
@@ -924,6 +924,8 @@ def forward(
 
         loss = None
         if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
             if self.config.problem_type is None:
                 if self.num_labels == 1:
                     self.config.problem_type = "regression"
@@ -1046,6 +1048,8 @@ def forward(
 
         loss = None
         if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
             loss_fct = CrossEntropyLoss()
             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))