Can we access the part of the network after patching and reshape?

[ysig]

Hi,

I would like to use a ViT style transformer but after the patching/reshaping is performed.
I was curious how I can access it easily from the current code.

Thanks!