Fix apply_rotary_emb() to have both tensors on the same device

awq/modules/fused/attn.py

@@ -36,7 +36,7 @@ def apply_rotary_emb(
     xk_ = torch.view_as_complex(
         xk.float().reshape(*xk.shape[:-1], 2, -1).transpose(-2, -1).contiguous()
     )
-    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_).to(xq_.device)
     xq_out = torch.view_as_real(xq_ * freqs_cis).transpose(-2, -1).flatten(3)
     xk_out = torch.view_as_real(xk_ * freqs_cis).transpose(-2, -1).flatten(3)
     return xq_out.type_as(xq), xk_out.type_as(xk)