I’m not sure if this is the best way to do it, but this works:
torch.cat([ torch.index_select(a, 1, i).unsqueeze(0) for a, i in zip(A, ind) ])

Essentially what this does is apply the regular index_select to each batch-element of A and ind and then concatenate the results together.