Use irange in a few places (#55325)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/55325
Test Plan: Sandcastle
Reviewed By: SciPioneer
Differential Revision: D27573006
fbshipit-source-id: 647b5da3901e92c23e95b2fe5e833e9081d72837
