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
    from art.estimators.certification.derandomized_smoothing import PyTorchDeRandomizedSmoothing

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

    try:
        cifar_data = fix_get_cifar10_data[0][:50]
        cifar_labels = fix_get_cifar10_data[1][:50]

        art_model = PyTorchDeRandomizedSmoothing(
            model="vit_small_patch16_224",
            loss=torch.nn.CrossEntropyLoss(),
            optimizer=torch.optim.SGD,
            optimizer_params={"lr": 0.01},
            input_shape=(3, 32, 32),
            nb_classes=10,
            ablation_size=4,
            load_pretrained=True,
            replace_last_layer=True,
            verbose=False,
        )

        scheduler = torch.optim.lr_scheduler.MultiStepLR(art_model.optimizer, milestones=[1], gamma=0.1)

        head = {
            "weight": torch.tensor(
                np.load(
                    os.path.join(
                        os.path.dirname(os.path.dirname(__file__)),
                        "../../utils/resources/models/certification/smooth_vit/head_weight.npy",
                    )
                )
            ).to(device),
            "bias": torch.tensor(
                np.load(
                    os.path.join(
                        os.path.dirname(os.path.dirname(__file__)),
                        "../../utils/resources/models/certification/smooth_vit/head_bias.npy",
                    )
                )
            ).to(device),
        }
        art_model.model.head.load_state_dict(head)

        art_model.fit(cifar_data, cifar_labels, nb_epochs=2, update_batchnorm=True, scheduler=scheduler)
        preds = art_model.predict(cifar_data)

        gt_preds = np.load(
            os.path.join(
                os.path.dirname(os.path.dirname(__file__)),
                "../../utils/resources/models/certification/smooth_vit/cumulative_predictions.npy",
            )
        )

        np.array_equal(preds, gt_preds)

    except ARTTestException as e:
        art_warning(e)

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