riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 02:02:21 +06:00
Code Issues Actions 33 Packages Projects Releases Wiki Activity

Cast logits to fp32 at the end of TF_T5 (#12332)

Browse Source 
This change enables tf.keras.mixed_precision with bf16
This commit is contained in:
Michal Szutenberg 2021-08-03 21:02:59 +02:00 committed by GitHub
parent b7439675b8
commit f064e0a43d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/models/t5/modeling_tf_t5.py
View File

@ -1407,6 +1407,8 @@ class TFT5ForConditionalGeneration(TFT5PreTrainedModel, TFCausalLanguageModeling
else:
logits = self.lm_head(sequence_output)
logits = tf.cast(logits, tf.float32)
loss = None if inputs["labels"] is None else self.compute_loss(inputs["labels"], logits)
if not inputs["return_dict"]: