Fix some convoluted(?) code (#48893)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48893
This simplifies some code which is written in an interesting way. It may be that this was intentional, but I don't recognize the pattern being used.
Test Plan: N/A - Sandcastle
Reviewed By: igorsugak
Differential Revision: D25358283
fbshipit-source-id: 19bcf01cbb117843e08df0237e6a03ea77958078
