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()