support Gather different indices for different examples in one batch (#23285)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/23285
for example:
Inputs:
data:
[[[2 4 2 0],
[0 1 2 0],
[1 1 0 0]],
[[3 4 1 3],
[0 3 2 2],
[4 1 0 4]]]
idx:
[[0 2],
[0 1]]
outputs:
[[[2 4 2 0],
[1 1 0 0]],
[[3 4 1 3],
[0 3 2 2]]]
data and idx must have the same outer dimension
call Gather or BatchGather with argument match_outer=True
Reviewed By: huayuli00
Differential Revision: D16452539
fbshipit-source-id: 7229489f4a9c02ee9f3c6a8a24bcd02925d96e07
