LR scheduler usage

1

How are you supposed to use an LR scheduler with SVI? The documentation for PyroLRScheduler currently has

optimizer = torch.optim.SGD
pyro_scheduler = pyro.optim.ExponentialLR({'optimizer': optimizer, 'optim_args': {'lr': 0.01}, 'gamma': 0.1})
svi = SVI(model, guide, pyro_scheduler, loss=TraceGraph_ELBO())
svi.step(),

which doesn’t show how to step the LR scheduler. I assume the right thing to do is something like this?

for e in range(n_epochs):
    pyro_scheduler(e)
    for mb in minibatches:
        svi.step(mb)

It would be very helpful to add a snippet clarifying this to the documentation.

2

the scheduler steps every svi step (this is handled internally). the scheduler functionality was changed in pytorch 1.1 instead of using set_epoch().