[FSDP] Address comments on previous PR (#85490)
Address follow ups on https://github.com/pytorch/pytorch/pull/85223/
Differential Revision: [D39740878](https://our.internmc.facebook.com/intern/diff/D39740878/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85490
Approved by: https://github.com/awgu
