def test_ablation(art_warning, fix_get_mnist_data, fix_get_cifar10_data):
    """
    Check that the ablation is being performed correctly
    """
    import torch

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    from art.estimators.certification.derandomized_smoothing.ablators.pytorch import ColumnAblatorPyTorch

    try:
        cifar_data = fix_get_cifar10_data[0]

        col_ablator = ColumnAblatorPyTorch(
            ablation_size=4,
            channels_first=True,
            to_reshape=False,  # do not upsample initially
            mode="ViT",
            original_shape=(3, 32, 32),
            output_shape=(3, 224, 224),
        )

        cifar_data = torch.from_numpy(cifar_data).to(device)
        # check that the ablation functioned when in the middle of the image
        ablated = col_ablator.forward(cifar_data, column_pos=10)

        assert ablated.shape[1] == 4
        assert torch.sum(ablated[:, :, :, 0:10]) == 0
        assert torch.sum(ablated[:, :, :, 10:14]) > 0
        assert torch.sum(ablated[:, :, :, 14:]) == 0

        # check that the ablation wraps when on the edge of the image
        ablated = col_ablator.forward(cifar_data, column_pos=30)

        assert ablated.shape[1] == 4
        assert torch.sum(ablated[:, :, :, 30:]) > 0
        assert torch.sum(ablated[:, :, :, 2:30]) == 0
        assert torch.sum(ablated[:, :, :, :2]) > 0

        # check that upsampling works as expected
        col_ablator = ColumnAblatorPyTorch(
            ablation_size=4,
            channels_first=True,
            to_reshape=True,
            mode="ViT",
            original_shape=(3, 32, 32),
            output_shape=(3, 224, 224),
        )

        ablated = col_ablator.forward(cifar_data, column_pos=10)

        assert ablated.shape[1] == 4
        assert torch.sum(ablated[:, :, :, : 10 * 7]) == 0
        assert torch.sum(ablated[:, :, :, 10 * 7 : 14 * 7]) > 0
        assert torch.sum(ablated[:, :, :, 14 * 7 :]) == 0

        # check that the ablation wraps when on the edge of the image
        ablated = col_ablator.forward(cifar_data, column_pos=30)

        assert ablated.shape[1] == 4
        assert torch.sum(ablated[:, :, :, 30 * 7 :]) > 0
        assert torch.sum(ablated[:, :, :, 2 * 7 : 30 * 7]) == 0
        assert torch.sum(ablated[:, :, :, : 2 * 7]) > 0

    except ARTTestException as e:
        art_warning(e)

    Check that the ablation is being performed correctly
    
tests/estimators/certification/test_vision_transformers.py