Update gemma.py

This commit is contained in:
Daniel Han-Chen 2024-02-26 14:38:27 +11:00
commit cf16811ab9

View file

@ -631,8 +631,9 @@ class FastGemmaModel(FastLlamaModel):
# Downcast RoPE embedding to correct data type
if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \
and (module.cos_cached.dtype != expected_dtype):
module.cos_cached = module.cos_cached.to(expected_dtype)
module.sin_cached = module.sin_cached.to(expected_dtype)
module.cos_cached = module.cos_cached.to(expected_dtype)
module.sin_cached = module.sin_cached.to(expected_dtype)
pass
pass
pass