Fix RecurrentGemma device_map (#30273)

Commit

1 year ago

Fix RecurrentGemma device_map (#30273) * Switch to non persistant buffer * fix device mismatch issue due to cache * style