TypeError: unbind(): argument 'input' (position 1) must be Tensor, not list

I have the following code that uses the Captum library:

import numpy as np
import torch.utils.data
import torchio as tio
from captum.attr import Saliency

from dcan.data_sets.dsets import LoesScoreDataset
from reprex.models import AlexNet3D

def normalize_array(array):
    new_array = (array - array.min()) / (array.max() - array.min())

    return new_array

net = AlexNet3D(4608)
model_save_location = '/home/models/loes_scoring_03.pt'

example_file = \
    '/home/feczk001/shared/projects/S1067_Loes/data/MNI-space_Loes_data/' \
image = tio.ScalarImage(example_file)

image_tensor = image.data

image_tensor = torch.unsqueeze(image_tensor, dim=0)
image_tensor = normalize_array(image_tensor)

with torch.no_grad():
    output = net(image_tensor)

print("Using existing trained model")

csv_data_file = "./data/MNI-space_Loes_data.csv"
testset = LoesScoreDataset(
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

dataiter = iter(testloader)
images, labels = next(dataiter)

ind = 3

input = images[ind].unsqueeze(0)
input.requires_grad = True


saliency = Saliency(net)
grads = saliency.attribute(input)

One the last line of code, I am getting this error:

Traceback (most recent call last):
  File "/home/miran045/reine097/projects/loes-scoring-explainability/src/my_captum_test.py", line 60, in <module>
    grads = saliency.attribute(input)
  File "/home/miran045/reine097/projects/loes-scoring-explainability/venv/lib/python3.6/site-packages/captum/log/__init__.py", line 42, in wrapper
    return func(*args, **kwargs)
  File "/home/miran045/reine097/projects/loes-scoring-explainability/venv/lib/python3.6/site-packages/captum/attr/_core/saliency.py", line 131, in attribute
    self.forward_func, inputs, target, additional_forward_args
  File "/home/miran045/reine097/projects/loes-scoring-explainability/venv/lib/python3.6/site-packages/captum/_utils/gradient.py", line 119, in compute_gradients
    grads = torch.autograd.grad(torch.unbind(outputs), inputs)
TypeError: unbind(): argument 'input' (position 1) must be Tensor, not list

The variable input on the last line of the code has the type torch.Tensor.

What am I doing wrong here? I realize this isn’t a stand-alone example, so please let me know if you want to know the values or types of any variables at run-time.

Based on the error message it seems an outputs object is computed inside saliency.attribute(input), so check if your model is returning a list while a tensor is expected.

I solved it by replacing this line:

saliency = Saliency(net)

by these lines:

def wrapped_model(inp):
    return net(inp)[0]

saliency = Saliency(wrapped_model)