|
@@ -524,6 +524,7 @@ class HHCache(Cache):
|
|
|
self.value_cache[layer_idx] = self.value_cache[layer_idx][mask].view(bsz, num_heads, -1, head_dim)
|
|
|
self.accumulated_attention_scores[layer_idx] = self.accumulated_attention_scores[layer_idx][mask].view(bsz, num_heads, -1)
|
|
|
|
|
|
+ pdb.set_trace()
|
|
|
|
|
|
def reorder_cache(self, beam_idx: torch.LongTensor):
|
|
|
"""Reorders the cache for beam search, given the selected beam indices."""
|