[enc-dec cache] fix bug in indexing (#32370)

sanchit-gandhi · web-flow · commit 9451a385261b · 2024-08-01T16:05:27.000+08:00
diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py
@@ -1034,7 +1034,7 @@ def __getitem__(self, layer_idx: int) -> List[Tuple[torch.Tensor]]:
                 self.self_attention_cache.key_cache[layer_idx],
                 self.self_attention_cache.value_cache[layer_idx],
                 self.cross_attention_cache.key_cache[layer_idx],
-                self.cross_attention_cache.key_cache[layer_idx],
+                self.cross_attention_cache.value_cache[layer_idx],
             )
         else:
             raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}")

Original file line number	Diff line number	Diff line change
`@@ -1034,7 +1034,7 @@ def __getitem__(self, layer_idx: int) -> List[Tuple[torch.Tensor]]:`
`1034`	`1034`	`self.self_attention_cache.key_cache[layer_idx],`
`1035`	`1035`	`self.self_attention_cache.value_cache[layer_idx],`
`1036`	`1036`	`self.cross_attention_cache.key_cache[layer_idx],`
`1037`		`- self.cross_attention_cache.key_cache[layer_idx],`
	`1037`	`+ self.cross_attention_cache.value_cache[layer_idx],`
`1038`	`1038`	`)`
`1039`	`1039`	`else:`
`1040`	`1040`	`raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}")`