Set proper return type (#42454)
Summary:
This function was always expecting to return a `size_t` value
Pull Request resolved: https://github.com/pytorch/pytorch/pull/42454
Reviewed By: ezyang
Differential Revision: D22993168
Pulled By: ailzhang
fbshipit-source-id: 044df8ce17983f04681bda8c30cd742920ef7b1e
