Fix heatmap issue