I have 2 model, one is a base model and another one which has a slightly different architecture from the base model which I called new model. So I am trying to load the weight from base to new model using load_state_dict but I keep getting unexpected key error on the new module part even though I already initialize the weight.
This is my old model
class SfsNetPipeline(nn.Module):
    """ SfSNet Pipeline
    """
    def __init__(self):
        super(SfsNetPipeline, self).__init__()

        self.conv_model            = baseFeaturesExtractions()
        self.normal_residual_model = NormalResidualBlock()
        self.normal_gen_model      = NormalGenerationNet()
        self.albedo_residual_model = AlbedoResidualBlock()
        self.albedo_gen_model      = AlbedoGenerationNet()
        self.light_estimator_model = LightEstimator()

    def get_face(self, sh, normal, albedo):
        shading = get_shading(normal, sh)
        recon   = reconstruct_image(shading, albedo)
        return recon

    def forward(self, face):
        # Following is training pipeline
        # 1. Pass Image from Conv Model to extract features

        out_features = self.conv_model(face)
        # 2 a. Pass Conv features through Normal Residual
        out_normal_features = self.normal_residual_model(out_features)
        # 2 b. Pass Conv features through Albedo Residual
        out_albedo_features = self.albedo_residual_model(out_features)
        # 3 a. Generate Normal
        predicted_normal = self.normal_gen_model(out_normal_features)
        # 3 b. Generate Albedo
        predicted_albedo = self.albedo_gen_model(out_albedo_features)

        all_features = torch.cat((out_features, out_normal_features, out_albedo_features), dim=1)
        # Predict SH
        predicted_sh = self.light_estimator_model(all_features)
        # 4. Generate shading
        out_shading = get_shading(predicted_normal, predicted_sh)
        out_recon = reconstruct_image(out_shading, predicted_albedo)
        #out_recon_color = self.color_model(out_recon)


        return predicted_normal, predicted_albedo, predicted_sh, out_shading, out_recon

And this is the new model (I comment the part where it is different)
class SfSRGBNetPipeline(nn.Module):
    """ SfSNet Pipeline
    """
    def __init__(self):
        super(SfSRGBNetPipeline, self).__init__()

        self.conv_model            = baseFeaturesExtractions()
        self.normal_residual_model = NormalResidualBlock()
        self.normal_gen_model      = NormalGenerationNet()
        self.albedo_residual_model = AlbedoResidualBlock()
        self.albedo_gen_model      = AlbedoGenerationNet()
        self.light_estimator_model = LightEstimatorNew()
        #punyaku
        #self.color_model = ColorCorrector()
        self.light_residual_model = LightFeatures()
        self.light_gen_model      = LightGenerationNet()

    def get_face(self, sh, normal, albedo):
        shading = get_shading(normal, sh)
        recon   = reconstruct_image(shading, albedo)
        return recon

    def forward(self, face):
        # Following is training pipeline
        # 1. Pass Image from Conv Model to extract features

        out_features = self.conv_model(face)
        # 2 a. Pass Conv features through Normal Residual
        out_normal_features = self.normal_residual_model(out_features)
        # 2 b. Pass Conv features through Albedo Residual
        out_albedo_features = self.albedo_residual_model(out_features)
        # 3 a. Generate Normal
        predicted_normal = self.normal_gen_model(out_normal_features)
        # 3 b. Generate Albedo
        predicted_albedo = self.albedo_gen_model(out_albedo_features)
        
        #Added new module here
        out_light_features = self.light_residual_model(out_features)
        predicted_correction = self.light_gen_model(out_light_features)
        # 3 c. Estimate lighting
        # First, concat conv, normal and albedo features over channels dimension
        all_features = torch.cat((out_light_features, out_normal_features, out_albedo_features), dim=1)
        # Predict SH
        predicted_sh = self.light_estimator_model(all_features)
        # 4. Generate shading
        out_shading = get_shading(predicted_normal, predicted_sh)
        # 5. Reconstruction of image
        out_recon = reconstruct_image(out_shading, predicted_albedo)
        #out_recon_color = self.color_model(out_recon)


        return predicted_normal, predicted_albedo, predicted_correction, predicted_sh, out_shading, out_recon

and this is how I load the model
sfs_net_ori      = SfsNetPipeline()
sfs_netku        = SfSRGBNetPipeline()

if use_cuda:
        sfs_net_model = nn.DataParallel(sfs_net_model)
        sfs_net_model = sfs_net_model.cuda()
        sfs_netku.apply(weights_init)        
        sfs_netku = nn.DataParallel(sfs_netku)
        sfs_netku = sfs_netku.cuda()
        checkpoint = torch.load(model_dir + 'sfs_net_model_9.pkl')
        sfs_net_pretrained_old = checkpoint['model_state_dict']
        sfs_net_pretrained_new = sfs_netku.state_dict()  
        sfs_netku.load_state_dict(sfs_net_pretrained_old)