code: def _test_tracin_xor_setup(self, tmpdir: str, use_gpu: bool = False):
      net = BasicLinearNet(2, 2, 1)
  
      state = OrderedDict(
          [
              (
                  "linear1.weight",
                  torch.Tensor(
                      [[-1.2956, -1.4465], [-0.3890, -0.7420]]
                  ),
              ),
              ("linear1.bias", torch.Tensor([1.2924, 0.0021])),
              ("linear2.weight", torch.Tensor([[-1.2013, 0.7174]])),
              ("linear2.bias", torch.Tensor([0.5880])),
          ]
      )
      net.load_state_dict(state)
      net_adjusted = (
          _wrap_model_in_dataparallel(net) if use_gpu else net
      )
  
      checkpoint_name = "-".join(["checkpoint", "class", "0" + ".pt"])
      torch.save(
          net_adjusted.state_dict(),
          os.path.join(tmpdir, checkpoint_name),
      )
  
      state = OrderedDict(
          [
              (
                  "linear1.weight",
                  torch.Tensor(
                      [[-1.3238, -1.4899], [-0.4544, -0.7448]]
                  ),
              ),
              ("linear1.bias", torch.Tensor([1.3185, -0.0317])),
              ("linear2.weight", torch.Tensor([[-1.2342, 0.7741]])),
              ("linear2.bias", torch.Tensor([0.6234])),
          ]
      )
      net.load_state_dict(state)
      net_adjusted = (
          _wrap_model_in_dataparallel(net) if use_gpu else net
      )
  
      checkpoint_name = "-".join(["checkpoint", "class", "1" + ".pt"])
      torch.save(
          net_adjusted.state_dict(),
          os.path.join(tmpdir, checkpoint_name),
      )
  
      state = OrderedDict(
          [
              (
                  "linear1.weight",
                  torch.Tensor(
                      [[-1.3546, -1.5288], [-0.5250, -0.7591]]
                  ),
              ),
              ("linear1.bias", torch.Tensor([1.3432, -0.0684])),
              ("linear2.weight", torch.Tensor([[-1.2490, 0.8534]])),
              ("linear2.bias", torch.Tensor([0.6749])),
          ]
      )
      net.load_state_dict(state)
      net_adjusted = (
          _wrap_model_in_dataparallel(net) if use_gpu else net
      )
  
      checkpoint_name = "-".join(["checkpoint", "class", "2" + ".pt"])
      torch.save(
          net_adjusted.state_dict(),
          os.path.join(tmpdir, checkpoint_name),
      )
  
      state = OrderedDict(
          [
              (
                  "linear1.weight",
                  torch.Tensor(
                      [[-1.4022, -1.5485], [-0.5688, -0.7607]]
                  ),
              ),
              ("linear1.bias", torch.Tensor([1.3740, -0.1571])),
              ("linear2.weight", torch.Tensor([[-1.3412, 0.9013]])),
              ("linear2.bias", torch.Tensor([0.6468])),
          ]
      )
      net.load_state_dict(state)
      net_adjusted = (
          _wrap_model_in_dataparallel(net) if use_gpu else net
      )
  
      checkpoint_name = "-".join(["checkpoint", "class", "3" + ".pt"])
      torch.save(
          net_adjusted.state_dict(),
          os.path.join(tmpdir, checkpoint_name),
      )
  
      state = OrderedDict(
          [
              (
                  "linear1.weight",
                  torch.Tensor(
                      [[-1.4464, -1.5890], [-0.6348, -0.7665]]
                  ),
              ),
              ("linear1.bias", torch.Tensor([1.3791, -0.2008])),
              ("linear2.weight", torch.Tensor([[-1.3818, 0.9586]])),
              ("linear2.bias", torch.Tensor([0.6954])),
          ]
      )
      net.load_state_dict(state)
      net_adjusted = (
          _wrap_model_in_dataparallel(net) if use_gpu else net
      )
  
      checkpoint_name = "-".join(["checkpoint", "class", "4" + ".pt"])
      torch.save(
          net_adjusted.state_dict(),
          os.path.join(tmpdir, checkpoint_name),
      )
  
      state = OrderedDict(
          [
              (
                  "linear1.weight",
                  torch.Tensor(
                      [[-1.5217, -1.6242], [-0.6644, -0.7842]]
                  ),
              ),
              ("linear1.bias", torch.Tensor([1.3500, -0.2418])),
              ("linear2.weight", torch.Tensor([[-1.4304, 0.9980]])),
              ("linear2.bias", torch.Tensor([0.7567])),
          ]
      )
      net.load_state_dict(state)
      net_adjusted = (
          _wrap_model_in_dataparallel(net) if use_gpu else net
      )
  
      checkpoint_name = "-".join(["checkpoint", "class", "5" + ".pt"])
      torch.save(
          net_adjusted.state_dict(),
          os.path.join(tmpdir, checkpoint_name),
      )
  
      state = OrderedDict(
          [
              (
                  "linear1.weight",
                  torch.Tensor(
                      [[-1.5551, -1.6631], [-0.7420, -0.8025]]
                  ),
              ),
              ("linear1.bias", torch.Tensor([1.3508, -0.2618])),
              ("linear2.weight", torch.Tensor([[-1.4272, 1.0772]])),
              ("linear2.bias", torch.Tensor([0.8427])),
          ]
      )
      net.load_state_dict(state)
      net_adjusted = (
          _wrap_model_in_dataparallel(net) if use_gpu else net
      )
  
      checkpoint_name = "-".join(["checkpoint", "class", "6" + ".pt"])
      torch.save(
          net_adjusted.state_dict(),
          os.path.join(tmpdir, checkpoint_name),
      )
  
      state = OrderedDict(
          [
              (
                  "linear1.weight",
                  torch.Tensor(
                      [[-1.5893, -1.6656], [-0.7863, -0.8369]]
                  ),
              ),
              ("linear1.bias", torch.Tensor([1.3949, -0.3215])),
              ("linear2.weight", torch.Tensor([[-1.4555, 1.1600]])),
              ("linear2.bias", torch.Tensor([0.8730])),
          ]
      )
      net.load_state_dict(state)
      net_adjusted = (
          _wrap_model_in_dataparallel(net) if use_gpu else net
      )
  
      checkpoint_name = "-".join(["checkpoint", "class", "7" + ".pt"])
      torch.save(
          net_adjusted.state_dict(),
          os.path.join(tmpdir, checkpoint_name),
      )
  
      dataset = BinaryDataset(use_gpu)
  
      return net_adjusted, dataset
  
file_path: captum/tests/influence/_core/test_tracin_xor.py
imports: import os
  import tempfile
  from collections import OrderedDict
  from typing import Callable, cast, Optional
  import torch
  import torch.nn
  import torch.nn.functional
  from captum.influence._core.tracincp import TracInCP
  from parameterized import parameterized
  from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
  from tests.influence._utils.common import _wrap_model_in_dataparallel, BasicLinearNet, BinaryDataset, build_test_name_func, DataInfluenceConstructor
label_desc: Use image node_id 1 for calling the TestTracInXOR obj's underlying member method code with example usage: obj._test_tracin_xor_setup(tmpdir, use_gpu) and returns: net_adjusted,  dataset