Updated assert to remove check on 3rd dim for MHA (#39402)
Summary:
## Description
* Updated assert statement to remove check on 3rd dimension (features) for keys and values in MultiheadAttention / Transform
* The feature dimension for keys and values can now be of different sizes
* Refer to https://github.com/pytorch/pytorch/issues/27623
Pull Request resolved: https://github.com/pytorch/pytorch/pull/39402
Reviewed By: zhangguanheng66
Differential Revision: D21841678
Pulled By: Nayef211
fbshipit-source-id: f0c9e5e0f33259ae2abb6bf9e7fb14e3aa9008eb
