Fix bug simplifying if-then-else when it can be removed (#44462)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/44462
Test Plan: Imported from OSS
Reviewed By: mruberry
Differential Revision: D23671157
Pulled By: bertmaher
fbshipit-source-id: b9b92ad0de1a7bd9bc1fcac390b542d885d0ca58
