Indexing a variable with a variable

2

You can use index_select:

import torch
from torch.autograd import Variable
x = Variable(torch.randn(3,3), requires_grad=True)
idx = Variable(torch.LongTensor([0,1]))
print(x.index_select(0, idx))

Note that the index variable (idx) can’t have requires_grad set to True. The variable being indexed (x) can have requires_grad=True.

http://pytorch.org/docs/tensors.html#torch.Tensor.index_select

8 Likes