def test_pytorch_training(art_warning, fix_get_mnist_data, fix_get_cifar10_data):
    """
    Check that the training loop for pytorch does not result in errors
    """
    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)

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

            self.conv1 = nn.Conv2d(in_channels=6, 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=1568, 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)

    for dataset, dataset_name in zip([fix_get_mnist_data, fix_get_cifar10_data], ["mnist", "cifar"]):
        if dataset_name == "mnist":
            ptc = SmallMNISTModel().to(device)
            input_shape = (1, 28, 28)
        else:
            ptc = SmallCIFARModel().to(device)
            input_shape = (3, 32, 32)

        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(ptc.parameters(), lr=0.01, momentum=0.9)
        try:
            for ablation_type in ["column", "row", "block"]:
                classifier = PyTorchDeRandomizedSmoothing(
                    model=ptc,
                    clip_values=(0, 1),
                    loss=criterion,
                    optimizer=optimizer,
                    input_shape=input_shape,
                    nb_classes=10,
                    ablation_type=ablation_type,
                    ablation_size=5,
                    threshold=0.3,
                    algorithm="levine2020",
                    logits=True,
                )
                classifier.fit(x=dataset[0], y=dataset[1], nb_epochs=1)
        except ARTTestException as e:
            art_warning(e)

    Check that the training loop for pytorch does not result in errors
    
tests/estimators/certification/test_derandomized_smoothing.py