Where is the code that visualizes by time step?

[HYER1N]

In the paper, there is a visualization figure for each time step, but it is not in the code. If it's in the code, please let me know which part it is. If it's not in the code, please let me know which part I should check.

[wenqsun]

I think that you can add a dictionary in the class AttentionStore to collect the cross-attention map for each step

and add the cross-attention map of each step to the dict in between_steps function

[daeunni]

@wenqsun Thanks! so much helpful