riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 02:02:21 +06:00
Code Issues Actions 23 Packages Projects Releases Wiki Activity

Fix autocast incompatibility in RecurrentGemma (#30832)

Browse Source
This commit is contained in:
Phillip Rust 2024-06-19 09:59:34 +02:00 committed by GitHub
parent b275a41005
commit 7c71b61dae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
View File

@ -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