How to use the captum.attr._core.neuron.neuron_integrated_gradients.NeuronIntegratedGradients function in captum

def _ig_input_test_assert(
        self,
        model,
        target_layer,
        test_input,
        test_neuron,
        expected_input_ig,
        additional_input=None,
    ):
        for internal_batch_size in [None, 1, 20]:
            grad = NeuronIntegratedGradients(model, target_layer)
            attributions = grad.attribute(
                test_input,
                test_neuron,
                n_steps=200,
                method="gausslegendre",
                additional_forward_args=additional_input,
                internal_batch_size=internal_batch_size,
            )
            if isinstance(expected_input_ig, tuple):
                for i in range(len(expected_input_ig)):
                    for j in range(attributions[i].shape[0]):
                        assertArraysAlmostEqual(
                            attributions[i][j].squeeze(0).tolist(),
                            expected_input_ig[i][j],
                            delta=0.1,
                        )

def _ig_matching_test_assert(self, model, output_layer, test_input, baseline=None):
        out = model(test_input)
        input_attrib = IntegratedGradients(model)
        ig_attrib = NeuronIntegratedGradients(model, output_layer)
        for i in range(out.shape[1]):
            ig_vals = input_attrib.attribute(test_input, target=i, baselines=baseline)
            neuron_ig_vals = ig_attrib.attribute(test_input, (i,), baselines=baseline)
            assertArraysAlmostEqual(
                ig_vals.reshape(-1).tolist(),
                neuron_ig_vals.reshape(-1).tolist(),
                delta=0.001,
            )
            self.assertEqual(neuron_ig_vals.shape, test_input.shape)

def test_simple_neuron_integrated_gradient(self):
        net = BasicModel_MultiLayer().cuda()
        inp = torch.tensor(
            [
                [0.0, 100.0, 0.0],
                [20.0, 100.0, 120.0],
                [30.0, 10.0, 0.0],
                [0.0, 0.0, 2.0],
            ]
        ).cuda()
        self._data_parallel_test_assert(
            NeuronIntegratedGradients,
            net,
            net.relu,
            alt_device_ids=True,
            inputs=inp,
            neuron_index=3,
        )

def test_multi_input_integrated_gradient(self):
        net = BasicModel_MultiLayer_MultiInput().cuda()
        inp1, inp2, inp3 = (
            10 * torch.randn(12, 3).cuda(),
            5 * torch.randn(12, 3).cuda(),
            2 * torch.randn(12, 3).cuda(),
        )
        self._data_parallel_test_assert(
            NeuronIntegratedGradients,
            net,
            net.model.relu,
            inputs=(inp1, inp2),
            additional_forward_args=(inp3, 5),
            neuron_index=(3,),
            test_batches=True,
        )