Captum usage with CTC loss function

I’m trying to use Captum’s IntegratedGradients for input attribution with a 1D convolution network which uses the CTC loss function. I have an input of:

(batchsize, 1, 4096)

And an output of:

(batchsize, 420, 5)

Where there is a sequence length of 420, and each part of the sequence has 5 possible classifications (or words). So there are 420 different classifications making this a multi-multi class classification.

However, I’m having trouble figuring out how to use the “target” argument in the IntegratedGradients.attribute function even with using a batchsize of 1.

I would like to get the input attribution for the argmax along the 2nd dimension (420) so I can see what portion of the input is used for each of the 420 classifiers.

Ultimately, I’d also like to do this for a batch size larger then one as well.

Any help or direction would be appreciated! Thanks!