pytorch

bdaf3226 - [FSDP] Ensure that customized non tensor optimizer state can be saved (#99214)

[FSDP] Ensure that customized non tensor optimizer state can be saved (#99214) The current logic does not actually handle all different non-tensor optimizer states correctly. This PR fixes the issue and adds a test. This PR will solve https://github.com/pytorch/pytorch/issues/99079 Differential Revision: [D45021331](https://our.internmc.facebook.com/intern/diff/D45021331/) Pull Request resolved: https://github.com/pytorch/pytorch/pull/99214 Approved by: https://github.com/awgu, https://github.com/awaelchli