You are saving the DataParallel module’s state_dict but loading to the non-DP module. If you call load_state_dict on the DataParallel module, it will work. Alternatively, when saving, use state_dict of the dp_module.module.state_dict().