Advanced indexing of 3D tensor?

I want to implement advanced indexing like the function showed in numpy:

input = np.random.randn(7, 8, 9)
print input[np.arange(7), [7, 7, 5, 3, 6, 7, 7], :].shape

But the pytorch code throws error:

input = torch.randn(7, 8, 9)
input[torch.range(7), torch.LongTensor([7, 7, 5, 3, 6, 7, 7]), :]

Is there any solution?

Try

input[torch.arange(0, 7).long(), torch.LongTensor([7, 7, 5, 3, 6, 7, 7]), :]
2 Likes