Fix template (#9840)

This commit is contained in:
Julien Plu
2021-01-27 13:40:30 +01:00
committed by GitHub
parent c7b7bd9963
commit bd701ab1a0
2 changed files with 8 additions and 4 deletions

View File

@@ -777,7 +777,9 @@ class TFBertMainLayer(tf.keras.layers.Layer):
# Since we are adding it to the raw scores before the softmax, this is
# effectively the same as removing these entirely.
extended_attention_mask = tf.cast(extended_attention_mask, dtype=embedding_output.dtype)
extended_attention_mask = tf.multiply(tf.subtract(1.0, extended_attention_mask), -10000.0)
one_cst = tf.constant(1.0, dtype=embedding_output.dtype)
ten_thousand_cst = tf.constant(-10000.0, dtype=embedding_output.dtype)
extended_attention_mask = tf.multiply(tf.subtract(one_cst, extended_attention_mask), ten_thousand_cst)
# Prepare head mask if needed
# 1.0 in head_mask indicate we keep the head