support for loading on cpu first when using DDP

in this example:

# Restore
model = ...
optimizer = ...
checkpoint = torch.load('amp_checkpoint.pt')

model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
amp.load_state_dict(checkpoint['amp'])

# Continue training
...

one cannot load the model to CPU first before restoring weight. this is necessary as a way to not have OOM with RAM when using huge models.

The flow that needs support is:

model.load_state(... using_cpu)
model.cuda(x)
model = amp.initialize(...)

model = DDP(model)

Start with the checkpoint restore flow using torch.load, load_state_dict, amp.initialize, and DDP as shown in the issue, then trace how model state is loaded and moved to CUDA. Done means the model can load weights on CPU first, move to CUDA, initialize AMP, and wrap with DDP without requiring excessive RAM.