Merge pull request #1229 from Feng0w0/wan_rope

[bugfix][NPU]:Fix bug that correctly obtains device type
This commit is contained in:
Zhongjie Duan
2026-02-04 13:26:00 +08:00
committed by GitHub
2 changed files with 2 additions and 2 deletions

View File

@@ -93,7 +93,7 @@ def rope_apply(x, freqs, num_heads):
x = rearrange(x, "b s (n d) -> b s n d", n=num_heads)
x_out = torch.view_as_complex(x.to(torch.float64).reshape(
x.shape[0], x.shape[1], x.shape[2], -1, 2))
freqs = freqs.to(torch.complex64) if freqs.device == "npu" else freqs
freqs = freqs.to(torch.complex64) if freqs.device.type == "npu" else freqs
x_out = torch.view_as_real(x_out * freqs).flatten(2)
return x_out.to(x.dtype)