Move _PartialWrapper to module scope (#59660)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59660
Context https://github.com/pytorch/pytorch/issues/57352
Test Plan: Pytorch CI tests
Reviewed By: vkuzo
Differential Revision: D28972991
fbshipit-source-id: efc9dd3e90e18e1cdf27d5ef0f168abd8169bc42
