Thank you. I have tried the following code snippet for the model according to the link you proposed 15 and another comment 15:
model = Model() #build the model
model = nn.DataParallel(model)
model.to(device) #move the model to device 
optimizer = optim.Adam(model.parameters()) #build the optimizer
try:
    model.load_state_dict(checkpoint['model'])
except RuntimeError:
    model.module.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])