I want to get neuron contributions within a layer of llama2 using captum’s LayerConductance and the following code:

I have pre-trained a "meta-llama/Llama-2-7b-chat-hf" model using the transformers library. Since my model uses additional tokens, I added them to the tokeniser before training and fine-tuned the “embed_tokens” module of the network. My training code looked like this: