ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update gemma.py

Browse source
This commit is contained in:
Daniel Han 2025-02-01 19:02:36 -08:00
commit 7ddeea39cc
1 changed files with 1 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

1
unsloth/models/gemma.py
View file

@ -223,6 +223,7 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module):
self.base = base
# Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this
self.current_rope_size = min(4 * 8192, self.max_position_embeddings)
print(dim, max_position_embeddings, base)
# Build here to make `torch.jit.trace` work.
self._set_cos_sin_cache(seq_len=self.current_rope_size, device=device, dtype=torch.get_default_dtype())