How to get the output to the original size for CViT

Hi,
I used 32*32 config for my case, I noticed the output file is 8,1024,1 where 8 is batch size, 1024 is the flatten result, and 1 is channel. if my input is 128 * 128, how to get back to the original size? is 1024 is equal to 32*32 (resize version of 128*128)? or is there something I must do to get back 128*128?

thank you
