another small fix for efficiency. Thanks stephen!

nod-ai · Nov 27, 2024 · b9d0f42 · b9d0f42
1 parent 3a474c4
commit b9d0f42
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/shortfin/python/shortfin_apps/llm/components/messages.py b/shortfin/python/shortfin_apps/llm/components/messages.py
@@ -58,8 +58,8 @@ def reset(self, phase: InferencePhase):
     def cache_page_indices(self, max_len: int) -> list[int]:
         if not self.allocation:
             return []
-        indices = [p.index for p in self.allocation.pages]
-        return indices[:max_len]
+        indices = [p.index for p in self.allocation.pages[:max_len]]
+        return indices
 
     def publish_allocated_pages(self, up_to_page_index: int):
         assert self.allocation