Add a condition for checking labels (#14211)

e823d819 · Haram Lee · GitHub · b3385963 · e823d819
Unverified Commit e823d819 authored Oct 30, 2021 by Haram Lee Committed by GitHub Oct 29, 2021
Show whitespace changes
Inline Side-by-side

Showing with 6 additions and 3 deletions

src/transformers/trainer_seq2seq.py src/transformers/trainer_seq2seq.py +6 -3

No files found.
--- a/src/transformers/trainer_seq2seq.py
+++ b/src/transformers/trainer_seq2seq.py
@@ -196,9 +196,12 @@ class Seq2SeqTrainer(Trainer):
        if self.args.prediction_loss_only:
            return (loss, None, None)

+        if has_labels:
            labels = inputs["labels"]
            if labels.shape[-1] < gen_kwargs["max_length"]:
                labels = self._pad_tensors_to_max_len(labels, gen_kwargs["max_length"])
+        else:
+            labels = None

        return (loss, generated_tokens, labels)