code: 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)
  
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 5 for calling a global function with example usage: test_pytorch_mnist_certification(art_warning, fix_get_mnist_data) and returns: self