About VAE in pyro

I have a dataset including large size images. In order to slove the memory limitations, I split images into 2500 patches.

Now I get latent space Z from the patch dataset, but a image still has 2500 * 32 dimension. I want to paste all Z of a image patches, and add a fully connected layer to reduce the latent space dim from 2500*32 into 20. How can I add this kind of linear layer into VAE and add loss into ELBO