set `tf.expand_dims argument` `axis` to an integer instead of a list in the multi_head_attention model.
@@ -424,10 +424,10 @@ class MultiHeadAttention(Layer):
if attention_mask is not None:
# The expand dim happens starting from the `num_heads` dimension,
# (<batch_dims>, num_heads, <query_attention_dims, key_attention_dims>)
- mask_expansion_axes = [-len(self._attention_axes) * 2 - 1]
+ mask_expansion_axis = -len(self._attention_axes) * 2 - 1
for _ in range(len(attention_scores.shape) - len(attention_mask.shape)):
attention_mask = tf.expand_dims(
- attention_mask, axis=mask_expansion_axes)
+ attention_mask, axis=mask_expansion_axis)
return self._softmax(attention_scores, attention_mask)
def _compute_attention(self,
