Generalise hard coded values

[rishabh-ranjan]

The code fails for general values of hidden dims, image sizes, etc. due to shape mismatches. Since image size and hidden dims are parameters anyway, can you please increase the flexibility in other parts of the code? This will allow it to be used as it is on other datasets and different architectures.

[PratikStar]

I can try to do this when I get some time :)

[clementchadebec]

Not sure it is still relevant to you but I have put together some Implementations of VAE variants that allow such flexibility here. Maybe this can help.