Problems using the cyclic scheduler

I was trying to follow the example and use a cyclic LR scheduler for the training and I set it up as follows:

optimizer =  torch.optim.SGD
scheduler = optim.CyclicLR({'optimizer': optimizer, 
                            'optim_args': {                            
                                'base_lr': 0.001, 
                                'max_lr': 1.0
                            }})

The base_lr and max_lr parameters are defined in the pytorch cyclic LR scheduler (CyclicLR — PyTorch 1.9.1 documentation)

Now, I use it as:

svi = SVI(my_model,
               my_guide,
               scheduler,
               loss=Trace_ELBO())

for i in range(num_iters):
    elbo = svi.step(frame)
    scheduler.step()

However, this results in the following error:

TypeError: __init__() got an unexpected keyword argument 'base_lr'

It seems it cannot wrap the underlying class parameters. What am I doing wrong here?

I believe optim_args are the arguments passed to the optimizer SGD (hence the error), so base_lr and max_lr, which are arguments to CyclicLR, should be moved up to the main argument dictionary:

optimizer =  torch.optim.SGD
scheduler = optim.CyclicLR({'optimizer': optimizer, 
                            'base_lr': 0.001,
                            'max_lr': 1.0,
                            'optim_args': {                            
                                'lr': 0.001,
                            }})