Revert "Avoid device casting for all singleton tensors in optimizer states (#91454)"
This reverts commit 1e725c97470d8cf74e85984ca997e77c76e91a18.
Reverted https://github.com/pytorch/pytorch/pull/91454 on behalf of https://github.com/janeyx99 due to Likely caused regression where checkpoint resume fails during training
