diff --git a/tiatoolbox/utils/visualization.py b/tiatoolbox/utils/visualization.py
index 9f0d65acd..5fb72021b 100644
--- a/tiatoolbox/utils/visualization.py
+++ b/tiatoolbox/utils/visualization.py
@@ -167,7 +167,7 @@ def overlay_prediction_mask(
     ax.imshow(overlay)
     ax.axis("off")
     # Generate colour bar
-    cbar = plt.colorbar(**colorbar_params)
+    cbar = plt.colorbar(**colorbar_params, ax=ax)
     cbar.ax.set_yticklabels(name_list)
     cbar.ax.tick_params(labelsize=12)
 
@@ -297,7 +297,7 @@ def overlay_probability_map(
     ax.imshow(overlay)
     ax.axis("off")
     # generate colour bar
-    cbar = plt.colorbar(**colorbar_params)
+    cbar = plt.colorbar(**colorbar_params, ax=ax)
     cbar.ax.tick_params(labelsize=12)
 
     return ax