Remove for-loop for getting number of elements in favour of abstraction (#60872)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60872
Test Plan: Sandcastle
Reviewed By: ngimel
Differential Revision: D29406199
fbshipit-source-id: ae49672cf1bb370d574d0c21231477bb17dea0ca
