Support nn.Modules with tuple output

pytorch / captum

Model interpretability and understanding for PyTorch

https://captum.ai

BSD 3-Clause "New" or "Revised" License

4.73k stars 476 forks source link

Support nn.Modules with tuple output #1215

Open 18jeffreyma opened 7 months ago

18jeffreyma commented 7 months ago

🚀 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 commented 6 months ago

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