Hi everyone,
I know that in order to load weights for CPU model which was saved during training by 1 GPU, we can use 2 lines below:
net = Model()  # My own architecture I define
model_path = "path/to/model.pkl"
state_dict = torch.load(model_path, map_location={"cuda:0" : "cpu"}
net.load_state_dict(state_dict)

However, when I train model on 2 GPUs using DataParallel to wrap my net model, then saving with
net = Model()
net = torch.nn.DataParallel(net)
net.cuda()

# Training
training()

# Saving
model_path = "path/to/model.pkl"
torch.save(net.state_dict(), model_path"

I load it back for later use like below:
net = Model()
net.load_state(torch.load(model_path, map_location={"cuda" : "cpu"})

It doesn’t map correctly the weights and keys in the state_dict. So what is the solution for it? Should I change the value for map_location to different one?
Thank you.