def forward(self, x):
    x = self.qkv(self.norm(x))
    x = rearrange(x, 'n s (h d) -> (n h) s d', h=self.nheads)
    q,k,v = torch.chunk(x, 3, dim=-1)
    s = (q@k.transpose(1,2))/self.scale
    x = s.softmax(dim=-1)@v
    x = rearrange(x, '(n h) s d -> n s (h d)', h=self.nheads)
    return self.proj(x)