[Feature Request] Compatibility with ToMe Cross Attention merging

[vidiotgameboss]

Currently as it is you'd get an error saying RuntimeError: shape '[x]' is invalid for input of size y
As it is this would force you to disable cross attention compression using ToMe only to generate a heatmap.