Fix A_log shape in KimiDeltaAttention (num_heads -> head_dim)

#144
by akkikiki - opened
Files changed (1) hide show
  1. modeling_kimi_linear.py +1 -1
modeling_kimi_linear.py CHANGED
@@ -518,7 +518,7 @@ class KimiDeltaAttention(nn.Module):
518
  )
519
 
520
  self.A_log = torch.nn.Parameter(torch.log(torch.empty(
521
- self.num_heads, dtype=torch.float32).uniform_(1, 16)))
522
 
523
  self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
524
  self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
 
518
  )
519
 
520
  self.A_log = torch.nn.Parameter(torch.log(torch.empty(
521
+ self.head_dim, dtype=torch.float32).uniform_(1, 16)))
522
 
523
  self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
524
  self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)