Fix Trainer checkpoint resume crashing on CPU with multiple processes (#49123)
* fix(trainer): resume from checkpoint on CPU with multiple processes
Trainer._load_optimizer_and_scheduler passes self.args.device as
map_location when world_size > 1, to load the optimizer state directly
on each GPU. On CPU (e.g. torchrun --nproc_per_node=2 with the gloo
backend), self.args.device is an indexed device like "cpu:0" rather
than plain "cpu" - but torch's CPU deserializer only understands the
plain "cpu" location tag, so resuming crashes on every rank with:
RuntimeError: don't know how to restore data location of
torch.storage.UntypedStorage (tagged with cpu:0)
The single-process case already avoids this by using "cpu" directly.
Extend that same fallback to the multi-process case whenever the
device is CPU, and only use the indexed device object when it's
actually a GPU device.
Fixes #49121
* Apply review: simplify map_location, drop the added test
Per SunMarc's review: the branch fits in one line, and the test isn't
needed.
Co-authored-by: Neev Modh <141932641+neevmodh@users.noreply.github.com>