shape error during deep supervision

[aoi-hill]

Hi... my model (inheriting from nn.Module) has a forward method implemented that returns multiple (>1) tensors, each of which is used to calculate one loss. On using it with IG, I am getting following error trace. The model works fine without IG integration, so I wanted to understand where is the problem within IG code.

Traceback (most recent call last):
File "/home/hrishikesh/Projects/lab/SCNet/Tmp/testing.py", line 65, in
attributions, approximation_error = ig.attribute(sample_input, target=0,return_convergence_delta=True)
File "/home/hrishikesh/tools/anaconda3/envs/scnet/lib/python3.10/site-packages/captum/log/init.py", line 42, in wrapper
return func(*args, **kwargs)
File "/home/hrishikesh/tools/anaconda3/envs/scnet/lib/python3.10/site-packages/captum/attr/_core/integrated_gradients.py", line 286, in attribute
attributions = self._attribute(
File "/home/hrishikesh/tools/anaconda3/envs/scnet/lib/python3.10/site-packages/captum/attr/_core/integrated_gradients.py", line 351, in _attribute
grads = self.gradient_func(
File "/home/hrishikesh/tools/anaconda3/envs/scnet/lib/python3.10/site-packages/captum/_utils/gradient.py", line 112, in compute_gradients
outputs = _run_forward(forward_fn, inputs, target_ind, additional_forward_args)
File "/home/hrishikesh/tools/anaconda3/envs/scnet/lib/python3.10/site-packages/captum/_utils/common.py", line 536, in _run_forward
return _select_targets(output, target)
File "/home/hrishikesh/tools/anaconda3/envs/scnet/lib/python3.10/site-packages/captum/_utils/common.py", line 580, in _select_targets
num_examples = output.shape[0]
AttributeError: 'tuple' object has no attribute 'shape'