Programming Assignment: Transformers Architecture with TensorFlow: Error: AssertionError('Wrong values in out')

I am having issues with the decoder layer and transformer, both give me this error. I checked all of the functions and they seem to be correct. Any suggestions?

I think this issue has been resolved.