Support nn.Modules with tuple output

[18jeffreyma]

🚀 Feature

Current usage assumes that mdules must output a pytorch tensor and not a tuple: many modules in transformers library return multiple outputs, making captum not work with them (e.g. LlamaForCausalLM). This would be very useful in supporting more captum usage with auto-regressive models like Llama2.

[josiehong]

I met the same issue, and the solution here gives a temporal implementation: #308