class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
    def forward(self, x, mask=None):
        batch, seq, d_model = x.size()
        q = self.q_proj(x).view(batch, seq, heads, d_k)
        k = self.k_proj(x).view(batch, seq, heads, d_k)
        v = self.v_proj(x).view(batch, seq, heads, d_k)
        scores = torch.matmul(q, k.transpose(-2, -1))
        scores = scores / math.sqrt(d_k)
        attn = F.softmax(scores, dim=-1)
        return torch.matmul(attn, v)