def test_deepspeed_custom_activation_checkpointing_params_forwarded():
    """Test that the activation checkpointing parameters get passed to `deepspeed.checkpointing.configure`
    correctly."""
    import deepspeed

    strategy = DeepSpeedStrategy(
        partition_activations=True,
        cpu_checkpointing=True,
        contiguous_memory_optimization=True,
        synchronize_checkpoint_boundary=True,
    )
    fabric = Fabric(
        strategy=strategy,
        precision="16-mixed",
        accelerator="cuda",
        devices=1,
    )
    fabric.launch()
    model = nn.Linear(3, 3)
    optimizer = torch.optim.Adam(model.parameters())

    with mock.patch("deepspeed.checkpointing.configure", wraps=deepspeed.checkpointing.configure) as configure:
        fabric.setup(model, optimizer)

    configure.assert_called_with(
        mpu_=None,
        partition_activations=True,
        contiguous_checkpointing=True,
        checkpoint_in_cpu=True,
        profile=None,
    )
Test that the activation checkpointing parameters get passed to `deepspeed.checkpointing.configure`
    correctly.
tests/tests_fabric/strategies/test_deepspeed_integration.py