I trained a model on 2 GPU’s using the DataParallel function, and saved the model state dict using:
For some reason when loading the dict into a model using:
model = models.segmentation.deeplabv3_resnet101(pretrained = True) model.load_state_dict(torch.load(PATH))
I get a key mismatch. I realized this was due to how I saved my model after using DataParallel, which in retrospect should have been:
Does anyone know a way to salvage this situation? i.e how I can load my state dict into a model, even though I saved the model differently from the documentation?
For sanity’s purposes, I also saved a model using
model.module.state_dict() while using
DataParallel, and again using
model.state_dict() using only a single GPU, in which both cases loaded the state_dict with no issues at all. Help?