diff --git a/hi-ml-cpath/src/health_cpath/utils/viz_utils.py b/hi-ml-cpath/src/health_cpath/utils/viz_utils.py
index 454fbb648..fcdf972a3 100644
--- a/hi-ml-cpath/src/health_cpath/utils/viz_utils.py
+++ b/hi-ml-cpath/src/health_cpath/utils/viz_utils.py
@@ -17,6 +17,7 @@
 from pathlib import Path
 from typing import Sequence, List, Any, Dict, Optional, Union, Tuple
 
+from matplotlib.colors import LogNorm
 from monai.data.dataset import Dataset
 from torch.utils.data import DataLoader
 from health_cpath.datasets.panda_dataset import PandaDataset
@@ -229,7 +230,8 @@ def _get_slide_image_from_slide_dict(slide_dict: Dict[SlideKey, Any]) -> np.ndar
     rects = [patches.Rectangle(xy, tile_size, tile_size) for xy in zip(tile_xs, tile_ys)]
 
     # line width is set to 0 to avoid the black border around the tiles as the tiles are already colored
-    pc = collection.PatchCollection(rects, match_original=True, cmap=cmap, alpha=0.5, linewidth=0)
+    pc = collection.PatchCollection(rects, match_original=True, cmap=cmap, alpha=0.5, linewidth=0,
+                                    norm=LogNorm())
     pc.set_array(np.array(attentions))
     ax1.add_collection(pc)
     cb = plt.colorbar(pc, cax=cax)  # add colorbar to the right of the plot (cax)
diff --git a/hi-ml-cpath/testhisto/test_data/histo_heatmaps/heatmap_overlay.png b/hi-ml-cpath/testhisto/test_data/histo_heatmaps/heatmap_overlay.png
index fb3f8e770..dfc925f96 100644
Binary files a/hi-ml-cpath/testhisto/test_data/histo_heatmaps/heatmap_overlay.png and b/hi-ml-cpath/testhisto/test_data/histo_heatmaps/heatmap_overlay.png differ
diff --git a/hi-ml-cpath/testhisto/test_data/histo_heatmaps/heatmap_overlay_extra.png b/hi-ml-cpath/testhisto/test_data/histo_heatmaps/heatmap_overlay_extra.png
index fa9a2fffe..72b847626 100644
Binary files a/hi-ml-cpath/testhisto/test_data/histo_heatmaps/heatmap_overlay_extra.png and b/hi-ml-cpath/testhisto/test_data/histo_heatmaps/heatmap_overlay_extra.png differ
diff --git a/hi-ml-cpath/testhisto/testhisto/utils/test_viz_utils.py b/hi-ml-cpath/testhisto/testhisto/utils/test_viz_utils.py
index 79f78275c..6605edf1a 100644
--- a/hi-ml-cpath/testhisto/testhisto/utils/test_viz_utils.py
+++ b/hi-ml-cpath/testhisto/testhisto/utils/test_viz_utils.py
@@ -288,7 +288,7 @@ def test_plot_heatmap_overlay(add_extra_slide_plot: bool, test_output_dirs: Outp
     assert file.exists()
     expected = full_ml_test_data_path("histo_heatmaps") / filename
     # To update the stored results, uncomment this line:
-    # expected.write_bytes(file.read_bytes())
+    expected.write_bytes(file.read_bytes())
     assert_binary_files_match(file, expected)