Have BertEncoder ingest keras_nlp.SelfAttentionMask

PiperOrigin-RevId: 330623395

Have BertEncoder ingest keras_nlp.SelfAttentionMask
PiperOrigin-RevId: 330623395
47b1e6aa · Zhenyu Tan · A. Unique TensorFlower · dbbed31d · 47b1e6aa
Commit 47b1e6aa authored Sep 08, 2020 by Zhenyu Tan Committed by A. Unique TensorFlower Sep 08, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

official/nlp/modeling/networks/bert_encoder.py official/nlp/modeling/networks/bert_encoder.py +1 -1

No files found.
--- a/official/nlp/modeling/networks/bert_encoder.py
+++ b/official/nlp/modeling/networks/bert_encoder.py
@@ -168,7 +168,7 @@ class BertEncoder(tf.keras.Model):

    self._transformer_layers = []
    data = embeddings
-    attention_mask = layers.SelfAttentionMask()([data, mask])
+    attention_mask = keras_nlp.SelfAttentionMask()(data, mask)
    encoder_outputs = []
    for i in range(num_layers):
      if i == num_layers - 1 and output_range is not None: