C5W4A1 - Transformers : UNQ_C6 - Decoder Layer call() function

I am facing the below error in my implementation of the call() function in the DecoderLayer Class. I am pretty positive on my implementation being correct. Please take a look at it.
Error:

My code for the call() function:

Also, I am really finding it difficult to finish the next 2 graded functions (DECODER & TRANSFORMER). It will be very helpful if you could improve the help documentations.

The order of the arguments in your mha2() layer is incorrect.

@TMosh thanks for the reply.