(max_seq_len, dim, theta=10000)
| 223 | |
| 224 | @amp.autocast(enabled=False, device_type="cuda") |
| 225 | def rope_params(max_seq_len, dim, theta=10000): |
| 226 | assert dim % 2 == 0 |
| 227 | freqs = torch.outer( |
| 228 | torch.arange(max_seq_len), |
| 229 | 1.0 / torch.pow(theta, |
| 230 | torch.arange(0, dim, 2).to(torch.float64).div(dim))) |
| 231 | freqs = torch.polar(torch.ones_like(freqs), freqs) |
| 232 | return freqs |
| 233 | |
| 234 | |
| 235 | class WanRMSNorm(nn.Module): |