Internal change

PiperOrigin-RevId: 346790357

Internal change
PiperOrigin-RevId: 346790357
03d1cda8 · A. Unique TensorFlower · 5fc4c351 · 03d1cda8
Commit 03d1cda8 authored Dec 10, 2020 by A. Unique TensorFlower
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

official/nlp/projects/bigbird/attention.py official/nlp/projects/bigbird/attention.py +1 -1

No files found.
--- a/official/nlp/projects/bigbird/attention.py
+++ b/official/nlp/projects/bigbird/attention.py
@@ -476,7 +476,7 @@ class BigBirdAttention(tf.keras.layers.MultiHeadAttention):

    attention_output = self._compute_attention(query, key, value,
                                               attention_mask)
-    attention_output.set_shape([None, None, self._num_heads, self._key_dim])
+    attention_output.set_shape([None, None, self._num_heads, self._value_dim])
    attention_output = self._output_dense(attention_output)
    return attention_output