Thanks for getting back! In my case, I don’t want my desired output to be fixed – I’d like it to match the dimensions of the input, which was why I wanted to keep track of the dimension before the adaptive pooling layer.
Keeping the arguments in __init__() the same (self) – so no adding an argument for input dimension or anything – what do  you suggest is the best way to write a flexible VAE that will output something with the same dimensions? Sorry if this is a really general question