Fix dtype issue in TF BART (#15178)

9a2dabae · Matt · GitHub · 0167edc8 · 9a2dabae
Unverified Commit 9a2dabae authored Jan 17, 2022 by Matt Committed by GitHub Jan 17, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/transformers/models/bart/modeling_tf_bart.py src/transformers/models/bart/modeling_tf_bart.py +1 -1

No files found.
--- a/src/transformers/models/bart/modeling_tf_bart.py
+++ b/src/transformers/models/bart/modeling_tf_bart.py
@@ -1383,7 +1383,7 @@ class TFBartForConditionalGeneration(TFBartPretrainedModel, TFCausalLanguageMode
        if inputs["labels"] is not None:
            inputs["labels"] = tf.where(
                inputs["labels"] == self.config.pad_token_id,
-                tf.fill(shape_list(inputs["labels"]), -100),
+                tf.cast(tf.fill(shape_list(inputs["labels"]), -100), inputs["labels"].dtype),
                inputs["labels"],
            )
            inputs["use_cache"] = False