Make sure that indices are available on all relevant devices

dottxt-ai · Oct 29, 2023 · c3421aa · c3421aa
1 parent 95a9d97
commit c3421aa
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/outlines/text/generate/sequence.py b/outlines/text/generate/sequence.py
@@ -226,7 +226,8 @@ def __call__(
 
             if past_key_values:
                 unfinished_past_key_values = tuple(
-                    tuple(vv[~local_is_finished] for vv in v) for v in past_key_values
+                    tuple(vv[~local_is_finished.to(vv.device)] for vv in v)
+                    for v in past_key_values
                 )
 
         result = self.model.tokenizer.decode(token_ids[:, num_prompt_tokens:])