Add handling of ifs to shape propagation (#62914)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/62914
Test Plan: Imported from OSS
Reviewed By: gchanan
Differential Revision: D30196945
Pulled By: eellison
fbshipit-source-id: 1c0c7f938c4547330fd1dba8ab7dd0b99a79b6a9
