code: 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)
  
file_path: adversarial-robustness-toolbox/tests/estimators/certification/test_vision_transformers.py
imports: import pytest
  import os
  import numpy
  from art.utils import load_dataset
  from tests.utils import ARTTestException
label_desc: Use image node_id 5 for calling a global function with example usage: test_pytorch_training(art_warning, fix_get_mnist_data, fix_get_cifar10_data) without return types