Revisit _move_optimizer_state function for all Strategies

[four4fish]

Proposed refactor

From comments in #10596
The device isn't used anymore

Motivation

Pitch

Additional context

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

cc @justusschock @awaelchli @akihironitta

[four4fish]

@tchaton seems the device still in use in GPUAccelerator's teardown function. Instead of move optimizer state to root_device, it defined move_optimizer_state to cpu to avoid memory leaking.
https://github.com/PyTorchLightning/pytorch-lightning/blame/master/pytorch_lightning/accelerators/gpu.py#L78-L80

[tchaton]

Hey @four4fish, sounds good to me. I believe the code can be shared across strategy though. Need to revisit the TPU one.