code: def test_pytorch_mnist_cnn(torch_device, interim):
      """The same test as above, but for pytorch"""
      torch = pytest.importorskip("torch")
  
      from torch import nn
      from torch.nn import functional as F
  
      class RandData:
          """Random test data."""
  
          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
  
      class Net(nn.Module):
          """Basic conv net."""
  
          def __init__(self):
              super().__init__()
              # Testing several different activations
              self.conv_layers = nn.Sequential(
                  nn.Conv2d(1, 10, kernel_size=5),
                  nn.MaxPool2d(2),
                  nn.Tanh(),
                  nn.Conv2d(10, 20, kernel_size=5),
                  nn.ConvTranspose2d(20, 20, 1),
                  nn.AdaptiveAvgPool2d(output_size=(4, 4)),
                  nn.Softplus(),
              )
              self.fc_layers = nn.Sequential(
                  nn.Linear(320, 50),
                  nn.BatchNorm1d(50),
                  nn.ReLU(),
                  nn.Linear(50, 10),
                  nn.ELU(),
                  nn.Softmax(dim=1),
              )
  
          def forward(self, x):
              """Run the model."""
              x = self.conv_layers(x)
              x = x.view(-1, 320)
              x = self.fc_layers(x)
              return x
  
      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.mse_loss(
                  output, torch.eye(10).to(device)[target]
              )
  
              loss.backward()
              optimizer.step()
  
              if num_examples > cutoff:
                  break
  
      # FIXME: this test should ideally pass with any random seed. See #2960
      random_seed = 42
  
      torch.manual_seed(random_seed)
      rs = np.random.RandomState(random_seed)
  
      batch_size = 32
  
      train_loader = RandData(batch_size)
      test_loader = RandData(batch_size)
  
      model = Net()
      optimizer = torch.optim.SGD(
          model.parameters(), lr=0.01, momentum=0.5
      )
  
      device = torch.device(torch_device)
  
      model.to(device)
      train(model, device, train_loader, optimizer, 1)
  
      next_x, _ = next(iter(train_loader))
      inds = rs.choice(next_x.shape[0], 3, replace=False)
  
      next_x_random_choices = next_x[inds, :, :, :].to(device)
  
      if interim:
          e = shap.DeepExplainer(
              (model, model.conv_layers[0]), next_x_random_choices
          )
      else:
          e = shap.DeepExplainer(model, next_x_random_choices)
  
      test_x, _ = next(iter(test_loader))
      input_tensor = test_x[:1].to(device)
      shap_values = e.shap_values(input_tensor)
  
      model.eval()
      model.zero_grad()
  
      with torch.no_grad():
          diff = (
              (model(input_tensor) - model(next_x_random_choices))
              .detach()
              .cpu()
              .numpy()
              .mean(0)
          )
  
      sums = np.array(
          [shap_values[i].sum() for i in range(len(shap_values))]
      )
      d = np.abs(sums - diff).sum() / np.abs(diff).sum()
  
      assert (
          d < 0.001
      ), f"Sum of SHAP values does not match difference! {d}"
  
file_path: shap/tests/explainers/test_deep.py
imports: import numpy
  import pandas
  import pytest
  from packaging import version
  import shap
  from shap import DeepExplainer
label_desc: Use image node_id 6 for calling a global function with example usage: test_pytorch_mnist_cnn(torch_device, interim) and returns: self, x, torch,  torch