[Tensor]#3 indexing with tensor

index with 1D array

it is possible to index a tensor using a different tensor (multiple values)

mask = tensor([1,0,0,0,0,0,0,0,...,0]) # only the 1st value is one
index = tensor([284, 149, 0, 1])
mask[index]
# tensor([0, 0, 1, 0], device='cuda:0')

mask[[284, 149, 0, 1]]
# error

index with matrix

Basically, indexing with tensor or ndarray is related to catch rows of targeted tensor or ndarray. Check the examples below.

A=torch.tensor([[0.5279, 0.5277, 0.0102],
            [0.9036, 0.9676, 0.2262],
            [0.9512, 0.6261, 0.2958]])
            
B=torch.tensor(2)
A.shape, B.shape, A[B] 
# (torch.Size([3, 3]), torch.Size([]), tensor([0.9512, 0.6261, 0.2958]))

B=torch.tensor([2])
A.shape, B.shape, A[B]
# (torch.Size([3, 3]), torch.Size([1]), tensor([[0.9512, 0.6261, 0.2958]]))

B=torch.tensor([2,1])
A.shape, B.shape, A[B]
# (torch.Size([3, 3]), torch.Size([2]),
# tensor([[0.9512, 0.6261, 0.2958],
#         [0.9036, 0.9676, 0.2262]]))

B=torch.tensor([[2,1],[0,0]])
A.shape, B.shape, A[B] 
# (torch.Size([3, 3]),
# torch.Size([2, 2]),
# tensor([[[0.9512, 0.6261, 0.2958],
#          [0.9036, 0.9676, 0.2262]],
#         [[0.5279, 0.5277, 0.0102],
#          [0.5279, 0.5277, 0.0102]]]))