Indexing in two dimensional PyTorch Tensor using another Tensor
Suppose that tensor A is defined as:

 1  2  3  4
 5  6  7  8
 9 10 11 12
13 14 15 16


I'm trying to extract a flat array out of this matrix by using another tensor as indices. For example, if the second tensor is defined as:

0
1
2
3


I want the result of the indexing to be 1-D tensor with the contents:

1
6
11
16


It doesn't seem to behave like NumPy; I've tried A[:, B] but it just throws an error for not being able to allocate an insane amount of memory and I've no idea why!

1st Approach: using torch.gather
torch.gather(A, 1, B.unsqueeze_(dim=1))

if you want one-dimensional vector, you can add squeeze to the end:
torch.gather(A, 1, B.unsqueeze_(dim=1)).squeeze_()

2nd Approach: using list comprehensions
You can use list comprehensions to select the items at specific indexes, then they can be concatenated using the torch.stack. An importat point here is that you should not use torch.tensor to create a new tensor from a list, if you do, you will break the chain (you cannot calculate gradient through that node):
torch.stack([A[i, B[i]] for i in range(A.size()[0])])

