Multichannel attention: Override _build_attention method instead of build()

PiperOrigin-RevId: 318208409

Multichannel attention: Override _build_attention method instead of build()
PiperOrigin-RevId: 318208409
658a45f2 · Hongkun Yu · A. Unique TensorFlower · 9e2050a6 · 658a45f2
Commit 658a45f2 authored Jun 24, 2020 by Hongkun Yu Committed by A. Unique TensorFlower Jun 24, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

official/nlp/modeling/layers/multi_channel_attention.py official/nlp/modeling/layers/multi_channel_attention.py +2 -2

No files found.
--- a/official/nlp/modeling/layers/multi_channel_attention.py
+++ b/official/nlp/modeling/layers/multi_channel_attention.py
@@ -117,8 +117,8 @@ class MultiChannelAttention(attention.MultiHeadAttention):
  cross-attention target sequences.
  """

-  def build(self, input_shape):
-    super(MultiChannelAttention, self).build(input_shape)
+  def _build_attention(self, qkv_rank):
+    super(MultiChannelAttention, self)._build_attention(qkv_rank)
    self._masked_softmax = masked_softmax.MaskedSoftmax(mask_expansion_axes=[2])

  def call(self, inputs, attention_mask=None):