Don't force dtype to torch.long in index_copy

PiperOrigin-RevId: 676937992
google-ai-edge · Sep 20, 2024 · 56abba0 · 56abba0
1 parent 44e57bb
commit 56abba0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ai_edge_torch/generative/layers/kv_cache.py b/ai_edge_torch/generative/layers/kv_cache.py
@@ -172,8 +172,8 @@ def _update_kv_base_impl(
     v_slice: torch.Tensor,
 ) -> KVCacheEntry:
   """Update the cache buffer without High Level Function Boundary annotation."""
-  k = cache.k_cache.index_copy(1, input_pos.to(torch.long), k_slice)
-  v = cache.v_cache.index_copy(1, input_pos.to(torch.long), v_slice)
+  k = cache.k_cache.index_copy(1, input_pos, k_slice)
+  v = cache.v_cache.index_copy(1, input_pos, v_slice)
   updated_cache = KVCacheEntry(k, v)
   return updated_cache