Proper map location for optimizer load
# What does this PR do?

I have been thinking more about #22159 and now remember why it might be better to load the optimizer state on the device directly: in multi-GPU training, the optimizer state is load in each process, so that would load it num_processes times on the CPU and risk a CPU RAM OOM.

Therefore, this adjusts #22159 to load the optimizer state:
- on CPU when there is only one process
- on each device directly when there are multiple.
_The documentation is not available anymore as the PR was closed or merged._