diff --git a/vis_utils/heatmap_utils.py b/vis_utils/heatmap_utils.py
index 1dc06d66..c8fc0e8b 100644
--- a/vis_utils/heatmap_utils.py
+++ b/vis_utils/heatmap_utils.py
@@ -20,7 +20,7 @@
 device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def score2percentile(score, ref):
-    percentile = percentileofscore(ref, score)
+    percentile = percentileofscore(ref.squeeze(), score.squeeze())
     return percentile
 
 def drawHeatmap(scores, coords, slide_path=None, wsi_object=None, vis_level = -1, **kwargs):