I am using DDP to distribute training across multiple gpu.
model = Net(...)
ddp_model = nn.SyncBatchNorm.convert_sync_batchnorm(model) 
ddp_model = DDP(ddp_model, device_ids=[gpu], find_unused_parameters=True)

When checkpointing, is it ok to save ddp_model.module instead of ddp_model? I need to be able to use the checkpoint for 1. evaluation using a single gpu 2. resume training with multiple gpus
def save_ckpt(ddp_model, path):
    state = {'model': ddp_model.module.state_dict(),
             'optimizer': optimizer.state_dict(),
            }
    torch.save(state, path)

def load_ckpt(path, distributed):
    checkpoint = torch.load(path, map_location=map_location)
    model = Net(...)
    optimizer = ...
    model.load_state_dict(checkpoint['model'], strict=False)
    optimizer.load_state_dict(checkpoint['optimizer'])
    if distributed:
        model = nn.SyncBatchNorm.convert_sync_batchnorm(model) 
        model = DDP(model, device_ids=[gpu], find_unused_parameters=True)
    return model

I am not sure how this would behave, does manipulating ddp_model.module in general break things?
Similarly what if the number of gpus changes, would that impact  the optimizer for example? does it need to be reinitialized? Thank you!