code: def test_pytorch_mnist_cnn():
      """The same test as above, but for pytorch"""
      # FIXME: this test should ideally pass with any random seed. See #2960
      random_seed = 0
  
      torch = pytest.importorskip("torch")
      torch.manual_seed(random_seed)
      rs = np.random.RandomState(random_seed)
  
      from torch import nn
      from torch.nn import functional as F
  
      batch_size = 128
  
      class RandData:
          """Ranomd data for testing."""
  
          def __init__(self, batch_size):
              self.current = 0
              self.batch_size = batch_size
  
          def __iter__(self):
              return self
  
          def __next__(self):
              self.current += 1
              if self.current < 10:
                  return torch.randn(
                      self.batch_size, 1, 28, 28
                  ), torch.randint(0, 9, (self.batch_size,))
              raise StopIteration
  
      try:
          # train_loader = torch.utils.data.DataLoader(
          #     datasets.MNIST(tmpdir, train=True, download=True,
          #                 transform=transforms.Compose([
          #                     transforms.ToTensor(),
          #                     transforms.Normalize((0.1307,), (0.3081,))
          #                 ])),
          #     batch_size=batch_size, shuffle=True)
          # test_loader = torch.utils.data.DataLoader(
          #     datasets.MNIST(tmpdir, train=False, download=True,
          #                 transform=transforms.Compose([
          #                     transforms.ToTensor(),
          #                     transforms.Normalize((0.1307,), (0.3081,))
          #                 ])),
          #     batch_size=batch_size, shuffle=True)
          train_loader = RandData(batch_size)
          test_loader = RandData(batch_size)
      except HTTPError:
          pytest.skip()
  
      def run_test(train_loader, test_loader, interim):
          class Net(nn.Module):
              """A test model."""
  
              def __init__(self):
                  super().__init__()
                  self.conv1 = nn.Conv2d(1, 5, kernel_size=5)
                  self.conv2 = nn.Conv2d(5, 10, kernel_size=5)
                  self.conv2_drop = nn.Dropout2d()
                  self.fc1 = nn.Linear(160, 20)
                  self.fc2 = nn.Linear(20, 10)
  
              def forward(self, x):
                  """Run the model."""
                  x = F.relu(F.max_pool2d(self.conv1(x), 2))
                  x = F.relu(
                      F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)
                  )
                  x = x.view(-1, 160)
                  x = F.relu(self.fc1(x))
                  x = F.dropout(x, training=self.training)
                  x = self.fc2(x)
                  return F.log_softmax(x, dim=1)
  
          model = Net()
          optimizer = torch.optim.SGD(
              model.parameters(), lr=0.01, momentum=0.5
          )
  
          def train(
              model, device, train_loader, optimizer, _, cutoff=20
          ):
              model.train()
              num_examples = 0
              for _, (data, target) in enumerate(train_loader):
                  num_examples += target.shape[0]
                  data, target = data.to(device), target.to(device)
                  optimizer.zero_grad()
                  output = model(data)
                  loss = F.nll_loss(output, target)
                  loss.backward()
                  optimizer.step()
                  # if batch_idx % 10 == 0:
                  #     print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                  #         epoch, batch_idx * len(data), len(train_loader.dataset),
                  #         100. * batch_idx / len(train_loader), loss.item()
                  #     ))
                  if num_examples > cutoff:
                      break
  
          device = torch.device("cpu")
          train(model, device, train_loader, optimizer, 1)
  
          next_x, _ = next(iter(train_loader))
          inds = rs.choice(next_x.shape[0], 3, replace=False)
          if interim:
              e = shap.GradientExplainer(
                  (model, model.conv1), next_x[inds, :, :, :]
              )
          else:
              e = shap.GradientExplainer(model, next_x[inds, :, :, :])
          test_x, _ = next(iter(test_loader))
          shap_values = e.shap_values(test_x[:1], nsamples=1000)
  
          if not interim:
              # unlike deepLIFT, Integrated Gradients aren't necessarily consistent for interim layers
              model.eval()
              model.zero_grad()
              with torch.no_grad():
                  diff = (
                      (model(test_x[:1]) - model(next_x[inds, :, :, :]))
                      .detach()
                      .numpy()
                      .mean(0)
                  )
              sums = np.array(
                  [
                      shap_values[i].sum()
                      for i in range(len(shap_values))
                  ]
              )
              d = np.abs(sums - diff).sum()
              assert d / (np.abs(diff).sum() + 0.01) < 0.1, (
                  "Sum of SHAP values "
                  "does not match difference! %f"
                  % (d / np.abs(diff).sum())
              )
  
      print("Running test from interim layer")
      run_test(train_loader, test_loader, True)
      print("Running test on whole model")
      run_test(train_loader, test_loader, False)
  
file_path: shap/tests/explainers/test_gradient.py
imports: from urllib.error import HTTPError
  import numpy
  import pandas
  import pytest
  import shap
label_desc: Use image node_id 2 for calling a global function with example usage: test_pytorch_mnist_cnn() and returns: self, torch,  torch, F