def test_optimizer_toggle_model_context_manager():
    """Test toggle_model runs when only one optimizer is used."""
    model = BoringModel()
    trainer = Mock()
    model.trainer = trainer
    params = model.parameters()
    optimizer = torch.optim.SGD(params, lr=0.1)
    trainer.optimizers = [optimizer]

    assert not model._param_requires_grad_state
    # toggle optimizer was failing with a single optimizer
    with model.toggled_optimizer(optimizer):
        assert model._param_requires_grad_state
    assert not model._param_requires_grad_state
Test toggle_model runs when only one optimizer is used.
tests/tests_pytorch/core/test_lightning_module.py