Load_model a model saved with DataParallelin pytorch


I like to use a pre-trained model in pytorch. Although I changed the device type to ‘cpu’, I still get the following error:
“RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cpu”
The model is originally saved with DataParallel.
I change the device of both data and model to ‘cpu’.

The test code is from:

You could try to save the state_dict of the underlying model without the nn.DataParallel wrapper as described here. Additionally, you could also push it to the CPU before saving it. This would allow you to load the model on any system later.

1 Like

Thanks, I could not use this solution, because the model is previously saved by someone else and I just want to test it. However, I found the solution by changing the model to the following:
self.model = self.model.module.cpu()
More info can be found in the following issue: