Checkpoinintg complex model

[jakubMitura14]

Hello, my model is created from two neural network one is Vnet which supplies segmentation masks second is modelRegression which count the class instances in the segmentation map.
both are kept in the lightning model class

        self.net = net
        self.modelRegression = UNetToRegresion(2,regression_channels)

However when I try the checkpointing the load_state_dict function gives error, what can I do?

trainer.lightning_module.load_state_dict(state_dict)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1604, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for Model:
        Missing key(s) in state_dict: "modelRegression.model.0.weight_fake_quant.min_vals", "modelRegression.model.0.weight_fake_quant.max_vals", "modelRegression.model.1.weight_fake_quant.min_vals", "modelRegression.model.1.weight_fake_quant.max_vals", "modelRegression.model.2.weight_fake_quant.min_vals", "modelRegression.model.2.weight_fake_quant.max_vals", "modelRegression.model.3.weight_fake_quant.min_vals", "modelRegression.model.3.weight_fake_quant.max_vals".