Use cub::DeviceSelect::UniqueByKey for EmbeddingBackward (#68376)
Summary:
https://github.com/NVIDIA/cub/pull/405 is still under review, API might change before it finally lands into cub 1.16, please wait for https://github.com/NVIDIA/cub/pull/405 before merging this. Tested locally and tests pass.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/68376
Reviewed By: bdhirsh
Differential Revision: D34706782
Pulled By: ngimel
fbshipit-source-id: a465d39bc24354d1047af1ee85be05a1de361c86
(cherry picked from commit 68a69bbc5093fd12b1fbfd561b3a10baf5d3e5ba)
