The softmax attention do not use a cosine similarity

[songlilucky]

Thanks for your work. according to your code, softmax attention do not use a cosine similarity,Did I get something wrong?

[songlilucky]

    # Dot product attention along cameras
    dot = self.scale * torch.einsum('b n Q d, b n K d -> b n Q K', q, k)
    dot = rearrange(dot, 'b n Q K -> b Q (n K)')
    att = dot.softmax(dim=-1)