class_docstr: This layer applies a multi-head self- or cross-attention as described in
  `Attention is all you need <https://arxiv.org/abs/1706.03762>`_ paper
  
  Args:
      embed_dim (int): :math:`C_{in}` from an expected input of size :math:`(N, S, C_{in})`
      num_heads (int): Number of heads in multi-head attention
      attn_dropout (Optional[float]): Attention dropout. Default: 0.0
      bias (Optional[bool]): Use bias or not. Default: ``True``
  
  Shape:
      - Input:
         - Query tensor (x_q) :math:`(N, S, C_{in})` where :math:`N` is batch size, :math:`S` is number of source tokens,
      and :math:`C_{in}` is input embedding dim
         - Optional Key-Value tensor (x_kv) :math:`(N, T, C_{in})` where :math:`T` is number of target tokens
      - Output: same shape as the input
code: def forward_pytorch(
      self,
      x_q: Tensor,
      x_kv: Optional[Tensor] = None,
      key_padding_mask: Optional[Tensor] = None,
      attn_mask: Optional[Tensor] = None,
  ) -> Tensor:
      out, _ = F.multi_head_attention_forward(
          query=x_q,
          key=x_kv if x_kv is not None else x_q,
          value=x_kv if x_kv is not None else x_q,
          embed_dim_to_check=self.embed_dim,
          num_heads=self.num_heads,
          in_proj_weight=torch.empty([0]),
          in_proj_bias=self.qkv_proj.bias,
          bias_k=None,
          bias_v=None,
          add_zero_attn=False,
          dropout_p=self.attn_dropout.p,
          out_proj_weight=self.out_proj.weight,
          out_proj_bias=self.out_proj.bias,
          training=self.training,
          key_padding_mask=key_padding_mask,
          need_weights=False,
          attn_mask=attn_mask,
          use_separate_proj_weight=True,
          q_proj_weight=self.qkv_proj.weight[: self.embed_dim, ...],
          k_proj_weight=self.qkv_proj.weight[
              self.embed_dim : 2 * self.embed_dim, ...
          ],
          v_proj_weight=self.qkv_proj.weight[2 * self.embed_dim :, ...],
      )
      return out
  
file_path: ml-cvnets/cvnets/layers/multi_head_attention.py
imports: from typing import Optional
  import torch
  from torch import Tensor, nn
  from torch.nn import functional
  from cvnets.layers.base_layer import BaseLayer
  from cvnets.layers.dropout import Dropout
  from cvnets.layers.linear_layer import LinearLayer
  from utils import logger
label_desc: Use image node_id 5 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward_pytorch(x_q, x_kv, key_padding_mask, attn_mask) and returns: out