Select tensor slice along a dimension based on index
I have a PyTorch tensor of the following shape: (100, 5, 100). I need to convert it into a tensor of shape (100, 100) by selecting from each row only one item in the second dimension, meaning that of those 5 elements I only need one, with its corresponding 100 elements.
To do this operation I have a second tensor of shape (100,) with the indices that specify which of those 5 items should be selected in each row.
Is there a simple way to perform this selection without having to mess with the dimensions too much?

Suppose tensor with indicies called idx and have shape (100,). Tensor with values called source. Then to select:
result = source[torch.arange(100), idx]

