Fix test to use proper condition. (#52216)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/52216
Test Plan: Imported from OSS
Reviewed By: pbelevich
Differential Revision: D26427506
Pulled By: ailzhang
fbshipit-source-id: ba4f2f66794cb2843926e5566eb4d25582f7fb2b
