Incorrect indexing for max pooling layers with DeepLIFT

[jmschrei]

🐛 Bug

Between captum 0.5.0 and 0.6.0 the arguments grad_input and grad_output in the DeepLIFT implementation switched from being a tuple of length 1 to just being a PyTorch tensor. Some of the non-linearity functions switched correctly, such as nonlinear:

0.5.0:

captum/captum/attr/_core/deep_lift.py

Line 956 in c2be437

abs(delta_in) < eps, new_grad_inp[0], grad_output[0] * delta_out / delta_in

Now: https://github.com/pytorch/captum/blob/master/captum/attr/_core/deep_lift.py#L879

But I don't think the max pooling function got switched over.

0.5.0:

captum/captum/attr/_core/deep_lift.py

Line 1118 in c2be437

abs(delta_in) < eps, grad_input[0], unpool_grad_out_delta / delta_in

Now: https://github.com/pytorch/captum/blob/master/captum/attr/_core/deep_lift.py#L1023

This causes the first index of the tensor grad_input to be used rather than the full thing. I'm not sure I understood how this happened because grad_output in the same function seems to have been changed correctly.

0.5.0:

captum/captum/attr/_core/deep_lift.py

Line 1075 in c2be437

grad_output_updated = grad_output[0]

Now: https://github.com/pytorch/captum/blob/master/captum/attr/_core/deep_lift.py#L1023

[jmschrei]

Would it be possible to get a confirmation that this is, indeed, a bug, and not a misunderstanding of the code? Thanks!