Update modeling_baichuan.py

modeling_baichuan.py

@@ -114,8 +114,8 @@ class RotaryEmbedding(torch.nn.Module):
         t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
         freqs = torch.outer(t, self.inv_freq)
         emb = torch.cat((freqs, freqs), dim=-1)
-        self.register_buffer("cos_cached", emb.cos().to(dtype)[None, None, :, :], persistent=False)
-        self.register_buffer("sin_cached", emb.sin().to(dtype)[None, None, :, :], persistent=False)
+        self.register_buffer("cos_cached", emb.cos().to(self.inv_freq.device)[None, None, :, :], persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(self.inv_freq.device)[None, None, :, :], persistent=False)
 
     def forward(self, x, seq_len):
         # x: [bs, num_attention_heads, seq_len, head_size]