def test_magnitude_pruner_morethanhalf_block_size(out_channels, block_size):
    """
    Test block sparsity when the block size is greater than half the number of channels
    """
    conv2d = torch.nn.Conv2d(
        in_channels=3,
        out_channels=out_channels,
        kernel_size=(3, 3),
        bias=False,
        groups=1,
    )

    weight_tensor = torch.rand_like(conv2d.weight)
    weight_tensor[weight_tensor == 0] = 1.0
    conv2d.weight.data = weight_tensor

    config = MagnitudePrunerConfig.from_dict(
        {
            "global_config": {
                "scheduler": {"update_steps": [1, 2]},
                "initial_sparsity": 0.0,
                "target_sparsity": 0.5,
                "block_size": block_size,
            }
        },
    )
    pruner = MagnitudePruner(conv2d, config)
    conv2d = pruner.prepare()

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

    if block_size > 1:
        block_sparse_channels = out_channels - out_channels % block_size
        for idx in range(0, block_sparse_channels, block_size):
            for jdx in range(1, block_size):
                assert torch.all(conv2d.weight_mask[idx] == conv2d.weight_mask[idx + jdx])

    sparsity = conv2d.weight_mask.eq(0).sum() / conv2d.weight_mask.numel()
    assert np.isclose(sparsity, 0.5, rtol=0.05)

    Test block sparsity when the block size is greater than half the number of channels
    
coremltools/test/optimize/torch/pruning/test_magnitude_pruner.py