def forward(
        self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor
    ) -> torch.Tensor:
        """Compute MHA(Q, K, V)

        Parameters
        ----------
        Q: torch.Tensor of size (batch_size, seq_len, input_dim)
            Input query (Q) sequence.
        K: torch.Tensor of size (batch_size, seq_len, input_dim)
            Input key (K) sequence.
        V: torch.Tensor of size (batch_size, seq_len, input_dim)
            Input value (V) sequence.

        Returns
        -------
        O: torch.Tensor of size (batch_size, seq_len, output_dim)
            Output MHA(Q, K, V)
        """
        assert Q.shape[-1] == K.shape[-1] == V.shape[-1] == self.input_dim

        batch_size, _, _ = Q.shape

        # embedding for multi-head projections (masked or not)
        Q = self.fc_q(Q)  # (B, S, D)
        K, V = self.fc_k(K), self.fc_v(V)  # (B, S, D)

        # Split into num_head vectors (num_heads * batch_size, n/m, head_dim)
        Q_ = torch.cat(Q.split(self.head_dim, -1), 0)  # (B', S, D')
        K_ = torch.cat(K.split(self.head_dim, -1), 0)  # (B', S, D')
        V_ = torch.cat(V.split(self.head_dim, -1), 0)  # (B', S, D')

        # Attention weights of size (num_heads * batch_size, n, m):
        # measures how similar each pair of Q and K is.
        W = torch.softmax(
            Q_.bmm(K_.transpose(-2, -1)) / math.sqrt(self.head_dim),
            -1,  # (B', D', S)
        )  # (B', N, M)

        # Multihead output (batch_size, seq_len, dim):
        # weighted sum of V where a value gets more weight if its corresponding
        # key has larger dot product with the query.
        H = torch.cat(
            (W.bmm(V_)).split(  # (B', S, S)  # (B', S, D')
                batch_size, 0
            ),  # [(B, S, D')] * num_heads
            -1,
        )  # (B, S, D)

        out = self.fc_o(H)

        return self.dropout(out)
Compute MHA(Q, K, V)

        Parameters
        ----------
        Q: torch.Tensor of size (batch_size, seq_len, input_dim)
            Input query (Q) sequence.
        K: torch.Tensor of size (batch_size, seq_len, input_dim)
            Input key (K) sequence.
        V: torch.Tensor of size (batch_size, seq_len, input_dim)
            Input value (V) sequence.

        Returns
        -------
        O: torch.Tensor of size (batch_size, seq_len, output_dim)
            Output MHA(Q, K, V)
        
braindecode/models/atcnet.py