def _attention(
    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Implementation of Scaled dot product attention"""
    # d_k - dimension of the query and key vectors
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    p_attn = F.softmax(scores, dim=-1)  # attention weights
    output = torch.matmul(p_attn, value)  # (B, h, T, d_k)
    return output, p_attn
Implementation of Scaled dot product attention
braindecode/models/sleep_stager_eldele_2021.py