Stage Heatmaps Not Calculated At All?

[viztopia]

Hi @elevanth! Thanks so much for your slim model code. I got the training up and running in my local environment, but the stage heat maps seem to be not calculated at all, after a day's training.. Below is a screenshot of the intermediate results. As you can see, the stage heatmaps are still completely dark. The gt_heatmap seems to be fine though..

Do you have any ideas about what might be causing this?