Minor bug fix to MultiHeadAttention registry function.
#456
+154
−39
MultiHeadAttention registry function.
#456