SA Output, Misalign w/ Fig. 4 Overview?

[GloryyrolG]

hi @omerbt @MichalGeyer @duongna21 @ all folks et. al.,

thx for ur contribution! may i ask if there is a mismatch between the extracted token is the output of self-attn in the paper & implementation in the code (seems to be latent, the output of the whole attn blk, ff, or the input of self-attn)? thx & best

TokenFlow/tokenflow_utils.py

Line 336 in 8ae24e9

self.pivot_hidden_states[0][batch_idxs].reshape(-1, dim))