Skip to content

Commit 6a26fc6

Browse files
committed
fix bug when doing beam search
Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>
1 parent 35cd0c1 commit 6a26fc6

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

optimum/exporters/ipex/cache_utils.py

+2
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,8 @@ def get_max_length(self) -> Optional[int]:
212212
def reset(self):
213213
"""Resets the cache values while preserving the objects"""
214214
self._seen_tokens = self.max_batch_size * [0]
215+
self.block_tables.fill_(-1)
216+
self.free_blocks = list(range(0, self.num_blocks))
215217

216218
def reorder_cache(self, beam_idx: torch.LongTensor):
217219
"""Reorders the cache for beam search, given the selected beam indices."""

0 commit comments

Comments
 (0)