Combing pytorches adam and scheduled learning rate into one (for when the model doesn't have a callback for the scheduler)

AdamStepLR.py

class AdamStepLR(torch.optim.Adam):
    """Combine Adam and lr_scheduler.StepLR so we can use it as a normal optimiser"""
    def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, step_size=50000, gamma=0.5):
        super().__init__(params, lr, betas, eps, weight_decay)
        self.scheduler = torch.optim.lr_scheduler.StepLR(self, step_size, gamma)
        
    def step(self):
        self.scheduler.step()
        return super().step()