Refactor ForeachUtils.h (#51131)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51131
--------
- Refactored `can_use_fast_route` logic in ForeachUtils.h.
- Fixed related bugs in test_foreach.py
Test Plan: Imported from OSS
Reviewed By: zou3519
Differential Revision: D26103904
Pulled By: izdeby
fbshipit-source-id: b3859b39adaab55c87dab6f7709d227adc0f6342
