def test_deepspeed_configure_optimizers():
    """Test that the deepspeed strategy with default initialization wraps the optimizer correctly."""
    from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer

    fabric = Fabric(
        strategy=DeepSpeedStrategy(),
        accelerator="cuda",
        devices=1,
        precision="16-mixed",
    )
    fabric.launch()
    model = nn.Linear(3, 3)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    model, optimizer = fabric.setup(model, optimizer)
    assert isinstance(optimizer.optimizer, DeepSpeedZeroOptimizer)
    assert isinstance(optimizer.optimizer.optimizer, torch.optim.SGD)
Test that the deepspeed strategy with default initialization wraps the optimizer correctly.
tests/tests_fabric/strategies/test_deepspeed_integration.py