def test_finalize(simple_module):
    """
    Test that calling finalize on the module leads to param being replaced with
    param_orig * param_mask.
    """
    config = MagnitudePrunerConfig.from_dict(
        {"global_config":
            {
                "scheduler": {"update_steps": [1, 2]},
                "initial_sparsity": 0.0,
                "target_sparsity": 0.5,
                "granularity": "per_channel"
            }},
    )
    pruner = MagnitudePruner(simple_module, config)
    simple_module = pruner.prepare()

    for _ in range(4):
        pruner.step()

    pruner.finalize(inplace=True)

    assert torch.sum(simple_module.conv2d.weight[:2] == 0).item() == 54
    assert torch.sum(simple_module.conv2d.weight[2] == 3).item() == 27
    assert torch.sum(simple_module.conv2d.weight[3] == 4).item() == 27

    Test that calling finalize on the module leads to param being replaced with
    param_orig * param_mask.
    
coremltools/test/optimize/torch/pruning/test_magnitude_pruner.py