Fix master - (#24003)
Summary:
I accidentally removed this in a merge, breaking a test. Fix for master
Pull Request resolved: https://github.com/pytorch/pytorch/pull/24003
Differential Revision: D16707108
Pulled By: eellison
fbshipit-source-id: 8b59f46e7932b88a7ae246a261c4daf17f23995f
