How can I make this PyTorch heatmap function faster and more efficient?

Question

I have this function that creates a sort if heatmap for 2d tensors, but it\'s painfully slow when using larger tensor inputs. How can I speed it up and make it more efficien