Ensure the query_states and key_states remain in bf16

#21

Applying the fix suggested from this thread.
Basically, making sure that after applying rotary embedding the data type does not change.

Microsoft org

LGTM. Thank you !

bapatra changed pull request status to merged

Sign up or log in to comment