Replace `internal::GRAIN_SIZE` by `grain_size` (parameter). (#53177)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/53013
Pull Request resolved: https://github.com/pytorch/pytorch/pull/53177
Reviewed By: SplitInfinity, nikithamalgifb
Differential Revision: D26860248
Pulled By: ngimel
fbshipit-source-id: 56917f8421f7b45c461945fd3d1ff107ce8535b2
