Fix autocast incompatibility in RecurrentGemma (#30832)

7c71b61d · Phillip Rust · GitHub · b275a410 · 7c71b61d
Unverified Commit 7c71b61d authored Jun 19, 2024 by Phillip Rust Committed by GitHub Jun 19, 2024
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py ...ormers/models/recurrent_gemma/modeling_recurrent_gemma.py +2 -2

No files found.
--- a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
+++ b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
@@ -254,8 +254,8 @@ class RecurrentGemmaSdpaAttention(nn.Module):
            k_out = k_out[:, :, indices]
            v_out = v_out[:, :, indices]

-            k_out[:, :, cache_position] = key_states
-            v_out[:, :, cache_position] = value_states
+            k_out[:, :, cache_position] = key_states.to(k_out.dtype)
+            v_out[:, :, cache_position] = value_states.to(v_out.dtype)

        self.key_states, self.value_states = k_out, v_out
        return k_out, v_out