# the output should be the same size as the input

Sample embeddings shape: torch.Size([128, 64, 256])
Output shape of mhsa: torch.Size([128, 64, 256])