Sorry, I should have included more information here. Every element in b is in [0, m). You can imagine that we have a sequence ranking task. k is the batch size, and for each batch instance, we have m candidate sequences of length n. So now b is the index of candidates, and there are k of them, because we want to choose one candidate for every batch instance. I think there should be some more straightforward way to handle this right? Since it kinda seems like a common need.