Get the value at a specific index in PyTorch

vision
1

I am stuck in this a bit for quite some time.

I have a ground truth label array for size 5.

y=tensor([903, 815, 214, 282, 694])

I have the output for scores array of shape : [5,1000]

    scores   =  tensor([[ 1.0406,  1.1808,  4.4227,  ...,  4.6864,  8.0145,  5.2128],
            [ 6.9101,  4.6083,  6.9259,  ...,  9.7415,  9.6305,  9.3974],
            [ 7.6097,  4.0396,  4.4560,  ...,  3.4892, 11.6411, 2],
            [ 1.0693,  4.6295,  5.3638,  ..., 10.9041, 10.8380,  9.2077],
            [ 1.7085,  1.4938,  8.6876,  ..., 15.1423,  9.6055,  9.8920]],
           grad_fn=<ViewBackward>)

I want the value from scores array based on the corresponding index of y. So for y[0], which is 903, I want the corresponding value from scores[1], position 903.

Is there some direct Pytorch function I can use.

2

Assuming you would like to get:

x[0, 903]
x[1, 815]
...

then this should work:

x = torch.randn(5, 1000)
idx = torch.tensor([903, 815, 214, 282, 694])

x[torch.arange(x.size(0)), idx]