Hi Alban (and smani)!



 albanD:

I think you want to expand your indices as Nx1xWxH by doing ind.view(N, 1, 1, 1).expand(N, 1, W, H).
Then you can use gather to get the values: res = value.gather(ind, 1).


How about direct indexing?
import torch
torch.__version__

torch.manual_seed (2020)

N = 5
C = 4
W = 3

t = torch.randn (N, C, W)
t

i = torch.LongTensor ([3, 0, 3, 1, 2])
i

t[torch.arange (0, N).long(), i, :]

>>> import torch
>>> torch.__version__
'0.3.0b0+591e73e'
>>>
>>> torch.manual_seed (2020)
<torch._C.Generator object at 0x00000274E0EF6630>
>>>
>>> N = 5
>>> C = 4
>>> W = 3
>>>
>>> t = torch.randn (N, C, W)
>>> t

(0 ,.,.) =
  1.2372 -0.9604  1.5415
 -0.4079  0.8806  0.0529
  0.0751  0.4777 -0.6759
 -2.1489 -1.1463 -0.2720

(1 ,.,.) =
  1.0066 -0.0416 -1.2853
 -0.4948 -1.2964 -1.2502
 -0.7693  1.6856  0.3546
  0.7790  0.3257  0.4995

(2 ,.,.) =
  0.7705 -0.5920  0.5270
  0.0807  0.9863  2.2251
 -1.1789 -0.3879 -0.3383
 -1.3321  0.2644 -0.4644

(3 ,.,.) =
  1.5745 -0.7463  0.5680
  0.5286  1.9171 -0.2539
 -1.2216 -0.0740 -1.8290
 -0.4015 -0.2322  0.2829

(4 ,.,.) =
  1.4925  0.7024 -1.5135
  0.1471  0.5943 -2.1462
  0.2443  0.4215 -0.4506
 -3.1644 -0.3315 -0.2427
[torch.FloatTensor of size 5x4x3]

>>>
>>> i = torch.LongTensor ([3, 0, 3, 1, 2])
>>> i

 3
 0
 3
 1
 2
[torch.LongTensor of size 5]

>>>
>>> t[torch.arange (0, N).long(), i, :]

-2.1489 -1.1463 -0.2720
 1.0066 -0.0416 -1.2853
-1.3321  0.2644 -0.4644
 0.5286  1.9171 -0.2539
 0.2443  0.4215 -0.4506
[torch.FloatTensor of size 5x3]

Best.
K. Frank