Eliminate narrowing conversion (#46730)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46730
A narrowing conversion on `last_idx` raises a compiler warning. This fixes that.
Test Plan: Standard pre-commit test rig.
Reviewed By: EscapeZero
Differential Revision: D24481497
fbshipit-source-id: f3e913b586738add59c422c3cf65035d87fc9e34