Update gemma.py

This commit is contained in:
Daniel Han 2025-02-01 19:11:20 -08:00
commit ff7cb20f93

View file

@ -211,6 +211,8 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module):
):
super().__init__()
if config is not None:
print(config)
print(dir(config))
# [TODO] Hack to pass in config - need to remove later
base = config.rope_theta
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0