Fix hardcoded heads in `BoxMultiHeadedAttention`

[deepsworld]

I confirm that this contribution is made under the terms of the license found in the root directory of this repository's source tree and that I have the authority necessary to make this contribution on behalf of its copyright owner.