code: def test_classification(self) -> None:
      def custom_baseline_fn(inputs: Tensor) -> Tensor:
          num_in = inputs.shape[1]  # type: ignore
          return torch.arange(0.0, num_in * 5.0).reshape(5, num_in)
  
      num_in = 40
      n_samples = 100
  
      # 10-class classification model
      model = SoftmaxModel(num_in, 20, 10)
      model.eval()
  
      inputs = torch.arange(0.0, num_in * 2.0).reshape(2, num_in)
      baselines = custom_baseline_fn
  
      self._assert_attributions(
          model, model.relu1, inputs, baselines, 1, n_samples
      )
  
file_path: captum/tests/attr/neuron/test_neuron_gradient_shap.py
imports: from typing import Callable, Tuple, Union
  import torch
  from captum.attr._core.neuron.neuron_gradient_shap import NeuronGradientShap
  from captum.attr._core.neuron.neuron_integrated_gradients import NeuronIntegratedGradients
  from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
  from tests.helpers.basic_models import BasicModel_MultiLayer
  from tests.helpers.classification_models import SoftmaxModel
  from torch import Tensor
  from torch.nn import Module
label_desc: Use image node_id 4 for calling the Test obj's underlying member method code with example usage: obj.test_classification() and returns: torch