Update benchmark.py

mr-mapache · Aug 8, 2024 · 4c86239 · 4c86239
1 parent 7a1360c
commit 4c86239
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/benchmark.py b/benchmark.py
@@ -125,8 +125,8 @@ def meta_forward(
 
     # Define parameters
     wq, wk, wv, wo = attention.q_projector, attention.k_projector, attention.v_projector, attention.output_projector
-    cache_k = attention.k_cache.sequence_cache.transpose(1, 2).detach().clone()
-    cache_v = attention.v_cache.sequence_cache.transpose(1, 2).detach().clone()
+    cache_k = attention.k_cache.sequence_cache.detach().clone()
+    cache_v = attention.v_cache.sequence_cache.detach().clone()
 
     n_local_heads = attention.number_of_heads
     n_local_kv_heads = attention.number_of_kv_heads