I use the nn.DataParallel to wrap my model to be trained on multiple GPUs and optimization was run successfully. Once the model was trained I save the model with the following code:
torch.save(cust_model.state_dict(), path)
Now when I want to make predictions, I use the following:
cust_model.load_state_dict(torch.load(model_dir, map_location=map_location_device))
, where the map_location_device is equal to “cpu” and model_dir is the model directory.
I get the following error when trying to load this model.
RuntimeError: Error(s) in loading state_dict for DenseSegmModel:
Missing key(s) in state_dict: “layer1.0.weight”, “layer1.1.weight”, “layer1.1.bias”, “layer2.denseblock1.denselayer1.norm1.weight”, “layer2.denseblock1.denselayer1.norm1.bias”,  and it goes on
Any ideas?