Select tensor slice along a dimension based on index

user_5331881 0
Viewed 151

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?

tensor pytorch python
1 Answers

Suppose tensor with indicies called idx and have shape (100,). Tensor with values called source. Then to select:

result = source[torch.arange(100), idx]
user_186356130
Related