Update gemma.py
This commit is contained in:
parent
911745377a
commit
c0f6b6c0cc
1 changed files with 1 additions and 1 deletions
|
|
@ -77,7 +77,7 @@ class FastGemmaRotaryEmbedding(torch.nn.Module):
|
|||
freqs = torch.outer(t, self.inv_freq)
|
||||
emb_old = torch.cat((freqs, freqs), dim=-1).to("cuda")
|
||||
|
||||
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
|
||||
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to("cuda")
|
||||
position_ids_expanded = position_ids[:, None, :].float()
|
||||
freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
|
||||
emb_new = torch.cat((freqs, freqs), dim=-1)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue