Minor bug fix to MultiHeadAttention registry function.#456

copybara-service[bot] wants to merge 1 commit intomastertensorflow/privacy:masterfrom

cl/520394717tensorflow/privacy:cl/520394717Copy head branch name to clipboard