code: def __init__(
      self,
      num_q_channels: int,
      num_kv_channels: int,
      num_heads: int,
      dropout: float,
  ):
      super().__init__()
      self.attention = nn.MultiheadAttention(
          embed_dim=num_q_channels,
          num_heads=num_heads,
          kdim=num_kv_channels,
          vdim=num_kv_channels,
          dropout=dropout,
          batch_first=True,
      )
  
imports: from torch import nn
label_desc: Use image node_id 1 to create a new MultiHeadAttention object from inherited base classes: nn with example: obj = MultiHeadAttention(num_q_channels, num_kv_channels, num_heads, dropout)