C4W2 Assignment NLP Transformer Summariser Error

Course Q&A Natural Language Processing NLP with Attention Models
10

I was not and I changed it to use tf.nn.softmax(scaled_attention_logits, axis=-1)

but unfortunately, I am still getting the same error message

InvalidArgumentError: Exception encountered when calling layer ‘softmax_58’ (type Softmax).

{{function_node _wrapped__AddV2_device/job:localhost/replica:0/task:0/device:GPU:0}} required broadcastable shapes [Op:AddV2] name:

Call arguments received by layer ‘softmax_58’ (type Softmax):
• inputs=tf.Tensor(shape=(1, 2, 2, 150), dtype=float32)
• mask=tf.Tensor(shape=(1, 1, 1, 2), dtype=float32)

Please help. Would it be possible for you to take a look at my notebook and see where I am going wrong?

Thank you very much.