Skip to content

Commit

Permalink
fix bug when doing beam search (#954)
Browse files Browse the repository at this point in the history
Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>
  • Loading branch information
kaixuanliu authored Oct 18, 2024
1 parent 35cd0c1 commit 80e8071
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions optimum/exporters/ipex/cache_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,8 @@ def get_max_length(self) -> Optional[int]:
def reset(self):
"""Resets the cache values while preserving the objects"""
self._seen_tokens = self.max_batch_size * [0]
self.block_tables.fill_(-1)
self.free_blocks = list(range(0, self.num_blocks))

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

0 comments on commit 80e8071

Please sign in to comment.