Update MultiHeadAttention docstring (#49950)

Summary:
Fixes MultiHeadAttention docstring.

Currently, https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html#torch.nn.MultiheadAttention
is

<img width="648" alt="Screen Shot 2020-12-29 at 21 06 43" src="https://user-images.githubusercontent.com/2459423/103311124-cd10cc00-4a19-11eb-89c9-0ee261364963.png">

and with the fix will be

<img width="648" alt="Screen Shot 2020-12-29 at 22 41 35" src="https://user-images.githubusercontent.com/2459423/103315838-0dc31200-4a27-11eb-82e2-ca8f13d713a1.png">

Pull Request resolved: https://github.com/pytorch/pytorch/pull/49950

Reviewed By: mrshenli

Differential Revision: D25732573

Pulled By: zhangguanheng66

fbshipit-source-id: b362f3f617ab26b0dd25c3a0a7d4117e522e620c
This commit is contained in:
vfdev-5 2021-01-05 13:28:48 -08:00 committed by Facebook GitHub Bot
parent 9945fd7253
commit e442ac1e3f

View File

@ -848,8 +848,9 @@ class MultiheadAttention(Module):
kdim: total number of features in key. Default: None.
vdim: total number of features in value. Default: None.
Note: if kdim and vdim are None, they will be set to embed_dim such that
query, key, and value have the same number of features.
Note that if :attr:`kdim` and :attr:`vdim` are None, they will be set
to :attr:`embed_dim` such that query, key, and value have the same
number of features.
Examples::