attn_weights = torch.softmax(attn_scores / d_out_kq**0.5, dim=1)
print(attn_weights)