PyTorch: how to select values from a multidimensional tensor using a multidimensional index?

Question

Given a tensor input of shape (5, 2, 2):

tensor([[[ 0,  1],
         [ 2,  3]],

        [[ 4,  5],
         [ 6,  7]],

        [[ 8,  9],
         [10, 11]],

        [[12, 13],
         [14, 15]],

        [[16, 17],
         [18, 19]]])

and a tensor index of shape (2,2):

tensor([[4, 3],
        [4, 2]])

How do I obtain the following output:

tensor([[16, 13],
        [18, 11]])

To put it into context: I have 5 input images of size 2x2 pixels (stored in input). Now I want to combine these 5 images into a single image of size 2x2 pixels, where index determines for each pixel from which input image it should be copied.

Example: starting at the top-left index[0,0] == 4, I take pixel value input[4,0,0] == 16. Then continue to index[0,1] == 3, I take pixel value input[3,0,1] == 13 and so on.

Answer 1

The solution I found is almost the same as this one:

indices = np.indices(input.shape)
indices[0] = index
input[tuple(indices)][0]

which gives the desired output

tensor([[16, 13],
        [18, 11]])