How to create a tensor by accessing specific values at given indices of a 2 X 2 tensor in pytorch?

Question

Suppose mat = torch.rand((5,7)) and I want to get values from 1st dimension (here, 7) by passing the indices, say idxs=[0,4,2,3,6]. The way I am able to do it now is by doing mat[[0,1,2,3,4],idxs]. I expected mat[:,idxs] to work, but it didn't. Is the first option the only way or is there a better way?

user_1714410 · Answer

torch.gather is what you are looking for:

torch.gather(mat, 1, torch.tensor(idxs)[:, None])

How to create a tensor by accessing specific values at given indices of a 2 X 2 tensor in pytorch?

1 Answers