Fix dimensionalities in apply_rotary_emb functions' comments

[tolgacangoz]

While working on MAGI-1, I came across these. I thought that understandability is important, so I propose fixing tensor shapes in the comments.

@yiyixuxu @a-r-r-o-w @ylacombe

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

thanks!

[tolgacangoz]

Thanks for merging!