Ensure the query_states and key_states remain in bf16 (#21)

- Ensure the query_states and key_states remain in bf16 (55b8e963ff0aa4a4190cff537165f08c378f62ff)


Co-authored-by: Mohammadreza Mohseni <[email protected]>

positional_embedding.py

@@ -269,10 +269,10 @@ class RotaryEmbedding(torch.nn.Module):
         return (
             apply_rotary_pos_emb(
                 q, cos_cached[seqlen_offset:seq_len], sin_cached[seqlen_offset:seq_len], seq_dimension=seq_dimension
-            ),
+            ).to(q.dtype),
             apply_rotary_pos_emb(
                 k, cos_cached[seqlen_offset:seq_len], sin_cached[seqlen_offset:seq_len], seq_dimension=seq_dimension
-            ),
+            ).to(k.dtype),
         )
 
     @classmethod