This is probably a very basic question but i’m having a hard time finding a solution.

Suppose I have a tensor of size 100 * 5 * 3 (batch_size * seq_len * distinct_class_probs)

I need to index dim=2 (distinct_class_probs) with a tensor.

So I want something like in each batch, for each token in a sequence, return the probability of a particular class. Which class probability to return will be specified by tensor. Tensor will be a one dimension tensor of size 100*5 (batch_size * seq_len)

Also the dimension of the expected return value would be 100 * 5

Any ideas on how i can achieve this? Thank you.