Pytorch Index a tensor of [n*n*3] by two tensors of size[n,n]
In numpy, I can index in the following:

a = np.random.randn(2,2,3)
b = np.eye(2,2).astype(np.uint8)
c = np.eye(2,2).astype(np.uint8)
print(a)
print("diff")
print(a[b,c,:])


, in which a[b, c, :] is a tensor of 2*2. 

[[[-1.01338087  0.70149058  0.55268617]
  [ 2.56941124  1.12720312 -0.07219555]]

 [[-0.04084548  0.17018995  2.14229567]
  [-0.68017558 -0.91788125  1.1719151 ]]]
diff
[[[-0.68017558 -0.91788125  1.1719151 ]
  [-1.01338087  0.70149058  0.55268617]]

 [[-1.01338087  0.70149058  0.55268617]
  [-0.68017558 -0.91788125  1.1719151 ]]]


But in Pytorch, I could not do the index in the same way like a[b,c,:]. Who knows how to do this. Thanks~


Indexing in PyTorch is almost similar to numpy.

a = torch.randn(2, 2, 3)
b = torch.eye(2, 2, dtype=torch.long)
c = torch.eye(2, 2, dtype=torch.long)

print(a)
print(a[b, c, :])




tensor([[[ 1.2471,  1.6571, -2.0504],
         [-1.7502,  0.5747, -0.3451]],

        [[-0.4389,  0.4482,  0.7294],
         [-1.3051,  0.6606, -0.6960]]])
tensor([[[-1.3051,  0.6606, -0.6960],
         [ 1.2471,  1.6571, -2.0504]],

        [[ 1.2471,  1.6571, -2.0504],
         [-1.3051,  0.6606, -0.6960]]])

