Add 64bit indexing support for softmax (#52713)
Summary:
fixes https://github.com/pytorch/pytorch/issues/52715 https://github.com/pytorch/pytorch/issues/52716
split across batch dimension
Pull Request resolved: https://github.com/pytorch/pytorch/pull/52713
Reviewed By: ailzhang
Differential Revision: D26640033
Pulled By: ngimel
fbshipit-source-id: f169cb0d6abc1cfbddf658d9775759a7d56f5c12
