In the tutorial on Variational Autoencoders (Variational Autoencoders — Pyro Tutorials 1.8.4 documentation), Trace_ELBO is used as the loss for svi optimization. I wonder if there is any more straightforward way of modifying Trace_ELBO such that the optimization can be done for beta-vae as well. It seems trivial to do that in pytorch since pytorch implement the loss of vae as reconstruction loss and regularization separately so the value of the beta parameter can be changed directly, but I’m not sure how to do it in pyro or if that is at all possible.
thank you so much! somehow I missed it