Ensure the query_states and key_states remain in bf16 (#21) f196467 verified bapatra mrm196 commited on Jun 12, 2024