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
      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)
  
file_path: adversarial-robustness-toolbox/tests/estimators/certification/test_derandomized_smoothing.py
imports: import os
  import pytest
  import numpy
  from art.utils import load_dataset
  from art.estimators.certification.derandomized_smoothing import PyTorchDeRandomizedSmoothing, TensorFlowV2DeRandomizedSmoothing
  from tests.utils import ARTTestException
label_desc: Use image node_id 3 for calling a global function with example usage: test_pytorch_training(art_warning, fix_get_mnist_data, fix_get_cifar10_data) and returns: self, self