Fix device placement of inv_freq after reset

trevor-m · Oct 7, 2024 · 1a7e62a · 1a7e62a
1 parent 2e11342
commit 1a7e62a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py
@@ -135,7 +135,7 @@ def _dynamic_frequency_update(self, position_ids, device):
             self.max_seq_len_cached = seq_len
 
         if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset
-            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
+            self.register_buffer("inv_freq", self.original_inv_freq.to(device), persistent=False)
             self.max_seq_len_cached = self.original_max_seq_len
 
     @torch.no_grad()