Update generator.py

Addressed the issue reported in #656 with the solution proposed by @rlouf . I am really not sure whether this is the best solution.
dottxt-ai · Feb 14, 2024 · c548a8a · c548a8a
1 parent b89e8df
commit c548a8a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/outlines/generate/generator.py b/outlines/generate/generator.py
@@ -265,7 +265,7 @@ def reorder_kv_cache(
     for cache_item in kv_cache:
         new_cache_item: Tuple = tuple()
         for layer in cache_item:
-            layer = torch.index_select(layer, 0, ancestors)
+            layer = torch.index_select(layer, 0, ancestors.to(layer.device))
             new_cache_item += (layer,)
         new_kv_cache += (new_cache_item,)