support Gather different indices for different examples in one batch (#23813)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/23813
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: D16652485
fbshipit-source-id: 9e144e97a8d6fceaf3b5714df1534338068f4a10
