Clone labels for end-dec model to solve RuntimeError

huggingface · Jan 10, 2024 · d153183 · d153183
1 parent 628e666
commit d153183
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py
@@ -662,6 +662,10 @@ def get_batch_logps(
         if not is_encoder_decoder:
             labels = labels[:, 1:].clone()
             logits = logits[:, :-1, :]
+        else:
+            # Fixes end-dec RuntimeError
+            labels = labels.clone()
+
         loss_mask = labels != label_pad_token_id
 
         # dummy token; we'll ignore the losses on these tokens later