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')
Ejemplo n.º 2
0
    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)