Batched index_select

[ratishsp (Ratish Puduppully)]

@jnhwkim

def batched_index_select(t, dim, inds): dummy = inds.unsqueeze(2).expand(inds.size(0), inds.size(1), t.size(2)) out = t.gather(dim, dummy) # b x e x f return out

Hi @jnhwkim,
The expected output is
(0 ,.,.) =
1 2
1 2

(1 ,.,.) =
5 6
5 6

(2 ,.,.) =
7 9
7 9
[torch.FloatTensor of size 3x2x2]

Running your version of batched_index_select results in a different output
(0 ,.,.) =
1 1 1
2 2 2

(1 ,.,.) =
5 5 5
6 6 6

(2 ,.,.) =
7 7 7
9 9 9
[torch.FloatTensor of size 3x2x3]