keys.shape: torch.Size([6, 2])
values.shape: torch.Size([6, 4])