MultiheadAttention
class torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None, batch_first=False, device=None, dtype=None)
Allows the model to jointly attend to information from different
   representation subspaces as described in the paper: Attention Is
   All You Need.
Multi-Head Attention is defined as:
  \text{MultiHead}(Q, K, V) =
  \text{Concat}(head_1,\dots,head_h)W^O

where head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V).
"forward()" will use a special optimized implementation if all of
   the following conditions are met:


self attention is being computed (i.e., "query", "key", and
     "value" are the same tensor. This restriction will be loosened in
     the future.)


inputs are batched (3D) with "batch_first==True"


Either autograd is disabled (using "torch.inference_mode" or
     "torch.no_grad") or no tensor argument "requires_grad"

