[FSDP] Code simplification
Pull Request resolved: https://github.com/pytorch/pytorch/pull/75227
Small change to reduce code duplication.
Differential Revision: [D35374608](https://our.internmc.facebook.com/intern/diff/D35374608/)
Approved by: https://github.com/zhaojuanmao
