fix nn.MHA scriptability (#58727)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58727
Test Plan: Imported from OSS
Reviewed By: ngimel
Differential Revision: D28593830
Pulled By: bhosmer
fbshipit-source-id: 37dee9efededaea9985a2bf040df1ba4b46f6580
