I found a function called [index_select](https://pytorch.org/docs/stable/torch.html#torch.index_select), but it doesn’t work for batch, because the argument index is a 1-D tensor. How can I do this for batch?