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

Update llama.py

Browse source
This commit is contained in:
Daniel Han 2024-07-23 11:16:31 -07:00
commit 7f5e3f0f82
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
unsloth/models/llama.py
View file

@ -1080,7 +1080,7 @@ class LlamaExtendedRotaryEmbedding(torch.nn.Module):
)
inv_freq = self.apply_scaling(inv_freq)
self.register_buffer("inv_freq", inv_freq, persistent = False)
print(__FILE__, __LINE__)
print(__LINE__)
# 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())
@ -1090,7 +1090,7 @@ class LlamaExtendedRotaryEmbedding(torch.nn.Module):
# Note: on the original Llama codebase, these tensors are created on the target device (and not on CPU) and
# in FP32. They are applied (multiplied) in FP32 as well.
self.current_rope_size = seq_len
print(__FILE__, __LINE__)
print(__LINE__)
t = torch.arange(self.current_rope_size, device="cpu", dtype=torch.int64).float()