diff --git a/README.md b/README.md
index df048f3..bc00b41 100644
--- a/README.md
+++ b/README.md
@@ -634,7 +634,7 @@ def display_qk_heatmap(qk_per_token):
     im = ax.imshow(qk_per_token.to(float).detach(), cmap='viridis')
     ax.set_xticks(range(len(prompt_split_as_tokens)))
     ax.set_yticks(range(len(prompt_split_as_tokens)))
-    ax.set_xticklabels(prompt_split_as_tokens)
+    ax.set_xticklabels(prompt_split_as_tokens, rotation=90)
     ax.set_yticklabels(prompt_split_as_tokens)
     ax.figure.colorbar(im, ax=ax)
     
diff --git a/images/implllama3_50_0.png b/images/implllama3_50_0.png
index c829d54..9da4b1b 100644
Binary files a/images/implllama3_50_0.png and b/images/implllama3_50_0.png differ
diff --git a/images/implllama3_52_0.png b/images/implllama3_52_0.png
index 78a4e8a..4a6309d 100644
Binary files a/images/implllama3_52_0.png and b/images/implllama3_52_0.png differ