Refactor to use shape arguments (#64299)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/64299
Test Plan: Imported from OSS
Reviewed By: driazati
Differential Revision: D30738141
Pulled By: eellison
fbshipit-source-id: 37ca30de81349ecf23d8656291863737b6ad6d96
