Move misplaced line (#29117)

Move misplaced line, improve code comment
huggingface · May 14, 2024 · fab977a · fab977a
1 parent de6b94c
commit fab977a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py
@@ -1176,11 +1176,11 @@ def forward(
  shift_logits = logits[..., :-1, :].contiguous()
  shift_labels = labels[..., 1:].contiguous()
  # Flatten the tokens
- loss_fct = CrossEntropyLoss()
  shift_logits = shift_logits.view(-1, self.config.vocab_size)
  shift_labels = shift_labels.view(-1)
- # Enable model parallelism
+ # Ensure tensors are on the same device
  shift_labels = shift_labels.to(shift_logits.device)
+ loss_fct = CrossEntropyLoss()
  loss = loss_fct(shift_logits, shift_labels)
 
  if not return_dict: