hmmm It could be, when you use data parallel state dict create a submodule called module. Try to save the state_dict instead of the model and before saving it allocate model in cpu and load state dict instead of that