fix indexing for more than 65535 elems in non-indexed first dim (#23123)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/22843, also adds test from https://github.com/pytorch/pytorch/issues/23102
Pull Request resolved: https://github.com/pytorch/pytorch/pull/23123
Differential Revision: D16402422
Pulled By: soumith
fbshipit-source-id: aa7a79159ed947be03ce3725ec8abcf5246a60bf
