I would like to make a neural network which uses black and white images as input and outputs a colored version of it. The important thing in that process is that the size of the images must stay the same. Usually this is done by using a Fully Convolutional Network with GAN or AE architecture. Now I have decided to implement a VAE version, but when i looked it up on the internet I found versions where the latent space was Linear/Dense meaning it breaks the Full Convolution. Is a VAE not a valid Approach for this type of neural network? Or is there a solution/some code someone can provide which would help in that situation? (Pytorch prefered)