Update fish_speech/models/text2semantic/llama.py

fish_speech/models/text2semantic/llama.py

@@ -1,3 +1,4 @@
+
 import dataclasses
 import json
 import math
@@ -862,7 +863,8 @@ class Attention(nn.Module):
 
         L, S = query.size(-2), key.size(-2)
         scale_factor = 1 / math.sqrt(query.size(-1))
-        attn_bias = torch.zeros(1, 1, L, S, dtype=query.dtype, device=query.device)
+        # FIX: Use new_zeros to avoid passing device object to torch.zeros which causes torch.compile error
+        attn_bias = query.new_zeros(1, 1, L, S)
 
         if attn_mask is not None:
             if attn_mask.dtype == torch.bool:
@@ -938,3 +940,4 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
 
     x_out2 = x_out2.flatten(3)
     return x_out2.type_as(x)
+