I was following the Pyro tutorial on stochastic variational inference and translating some things to NumPyro. Part IV: Tips and Tricks mentions to “10. Consider normalizing your ELBO” and references the Pyro tutorial on Scaling the Loss by using the
@poutine.scale(...) decorator. However, I couldn’t find how to do a similar scaling of model and guide functions in NumPyro.
How could I scale the (ELBO) loss in NumPyro?