def test_lambda_lr(self, lr_mock): state = {torchbearer.OPTIMIZER: 'optimizer'} scheduler = CosineAnnealingLR(4, eta_min=10, last_epoch=-4, step_on_batch='batch') scheduler.on_start(state) lr_mock.assert_called_once_with('optimizer', 4, eta_min=10, last_epoch=-4) self.assertTrue(scheduler._step_on_batch == 'batch')
def test_lambda_lr(self, lr_mock): state = {torchbearer.OPTIMIZER: 'optimizer', torchbearer.EPOCH: 0} scheduler = CosineAnnealingLR(T_max=4, eta_min=10, step_on_batch=True) scheduler.on_start(state) lr_mock.assert_called_once_with('optimizer', T_max=4, eta_min=10, last_epoch=-1) self.assertTrue(scheduler._step_on_batch)