inp will probably contain the prediction probabilities as [batch_size, nb_words].
If that’s the case, the gather operation will get all predictions for the current target word index in each sample:
batch_size = 5
nb_words = 10

inp = torch.randn(batch_size, nb_words)
target = torch.randint(0, nb_words, (batch_size,))
torch.gather(inp, 1, target.view(-1, 1))