def __init__(
        self,
        input_dim: int,
        head_dim: int,
        output_dim: int,
        num_heads: int,
        dropout: float = 0.0,
    ):
        """Multi-head Attention

        The difference between this module and torch.nn.MultiheadAttention is
        that this module supports embedding dimensions different then input
        and output ones. It also does not support sequences of different
        length.

        Parameters
        ----------
        input_dim : int
            Dimension of query, key and value inputs.
        head_dim : int
            Dimension of embed query, key and value in each head,
            before computing attention.
        output_dim : int
            Output dimension.
        num_heads : int
            Number of heads in the multi-head architecture.
        dropout : float, optional
            Dropout probability on output weights. Default: 0.0 (no dropout).
        """

        super(_MHA, self).__init__()

        self.input_dim = input_dim
        self.head_dim = head_dim
        # typical choice for the split dimension of the heads
        self.embed_dim = head_dim * num_heads

        # embeddings for multi-head projections
        self.fc_q = nn.Linear(input_dim, self.embed_dim)
        self.fc_k = nn.Linear(input_dim, self.embed_dim)
        self.fc_v = nn.Linear(input_dim, self.embed_dim)

        # output mapping
        self.fc_o = nn.Linear(self.embed_dim, output_dim)

        # dropout
        self.dropout = nn.Dropout(dropout)
Multi-head Attention

        The difference between this module and torch.nn.MultiheadAttention is
        that this module supports embedding dimensions different then input
        and output ones. It also does not support sequences of different
        length.

        Parameters
        ----------
        input_dim : int
            Dimension of query, key and value inputs.
        head_dim : int
            Dimension of embed query, key and value in each head,
            before computing attention.
        output_dim : int
            Output dimension.
        num_heads : int
            Number of heads in the multi-head architecture.
        dropout : float, optional
            Dropout probability on output weights. Default: 0.0 (no dropout).
        
braindecode/models/atcnet.py