def test_pytorch_mnist_certification(art_warning, fix_get_mnist_data):
    """
    Assert that the correct number of certifications are given for the MNIST dataset
    """
    import torch
    import torch.optim as optim
    import torch.nn as nn

    device = "cuda" if torch.cuda.is_available() else "cpu"

    class SmallMNISTModel(nn.Module):
        def __init__(self):
            super(SmallMNISTModel, self).__init__()

            self.conv1 = nn.Conv2d(in_channels=2, out_channels=32, kernel_size=(4, 4), dilation=(1, 1), stride=(2, 2))
            self.max_pool = nn.MaxPool2d(2, stride=2)
            self.fc1 = nn.Linear(in_features=1152, out_features=100)

            self.fc2 = nn.Linear(in_features=100, out_features=10)

            self.relu = nn.ReLU()

        def forward(self, x):
            if isinstance(x, np.ndarray):
                x = torch.from_numpy(x).to(device)
            x = self.relu(self.conv1(x))
            x = self.max_pool(x)
            x = torch.flatten(x, 1)
            x = self.relu(self.fc1(x))
            return self.fc2(x)

        def load_weights(self):
            fpath = os.path.join(
                os.path.dirname(os.path.dirname(__file__)), "../../utils/resources/models/certification/derandomized/"
            )
            self.conv1.weight = nn.Parameter(torch.from_numpy(np.load(fpath + "W_CONV2D1_MNIST.npy")).float())
            self.conv1.bias = nn.Parameter(torch.from_numpy(np.load(fpath + "B_CONV2D1_MNIST.npy")).float())

            self.fc1.weight = nn.Parameter(torch.from_numpy(np.load(fpath + "W_DENSE1_MNIST.npy")).float())
            self.fc1.bias = nn.Parameter(torch.from_numpy(np.load(fpath + "B_DENSE1_MNIST.npy")).float())

            self.fc2.weight = nn.Parameter(torch.from_numpy(np.load(fpath + "W_DENSE2_MNIST.npy")).float())
            self.fc2.bias = nn.Parameter(torch.from_numpy(np.load(fpath + "B_DENSE2_MNIST.npy")).float())

    ptc = SmallMNISTModel()
    ptc.load_weights()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(ptc.parameters(), lr=0.01)

    try:
        for ablation_type in ["column", "block"]:
            if ablation_type == "column":
                size_to_certify = 5
                ablation_size = 2
            else:
                """
                the model was trained on column ablations, so make the block task simpler so that a
                degree of certification is obtained.
                """
                size_to_certify = 1
                ablation_size = 5

            classifier = PyTorchDeRandomizedSmoothing(
                model=ptc,
                clip_values=(0, 1),
                loss=criterion,
                optimizer=optimizer,
                input_shape=(1, 28, 28),
                nb_classes=10,
                ablation_type=ablation_type,
                ablation_size=ablation_size,
                threshold=0.3,
                algorithm="levine2020",
                logits=True,
            )

            preds = classifier.predict(np.copy(fix_get_mnist_data[0]))
            cert, cert_and_correct, top_predicted_class_argmax = classifier.ablator.certify(
                preds, label=fix_get_mnist_data[1], size_to_certify=size_to_certify
            )
            if ablation_type == "column":
                assert np.sum(cert.cpu().numpy()) == 52
            else:
                assert np.sum(cert.cpu().numpy()) == 22
    except ARTTestException as e:
        art_warning(e)

    Assert that the correct number of certifications are given for the MNIST dataset
    
tests/estimators/certification/test_derandomized_smoothing.py