Update modeling_codegen.py

Hello, this small change ensures that the labels are on the correct device.

modeling_codegen.py

@@ -713,6 +713,7 @@ class CodeGenForCausalLM(CodeGenPreTrainedModel):
 
  loss = None
  if labels is not None:
+ labels = labels.to(lm_logits.device)
  # Shift so that tokens < n predict n
  shift_logits = lm_logits[..., :-1, :].contiguous()
  shift_labels = labels[..., 1:].contiguous()