Not able to save a new part of the model with save_checkpoint

[alessiabertugli]

Hi,

I have to load a pre-trained model and add a part of the model to complete the training process. The model is an autoencoder composed of an encoder, a latent code and a decoder. I have to add a new decoder that is initialised as a copy of the original decoder. I can train the whole model with both the decoders, but when I tried to save the checkpoint using the function save_checkpoint, the state_dict does not contain the new decoder parameters. I added the new decoder as follow:

model.new_decoder = copy.deepcopy(model.decoder)

I debugged the checkpoint_connector.py from pytorch-lightning and I found that the model contains the new decoder, but its parameters are not present calling the state_dict function as follow:

state_dict  = model.state_dict()

Looking at the state_dict function (from /torch/nn/modules/module.py) I see that self._modules contains the parameters of the new decoder but I can't understand why they are not put into the destination dictionary. Can anyone help me with this issue, please?

def state_dict(self, destination=None, prefix='', keep_vars=False):
        r"""Returns a dictionary containing a whole state of the module.

        Both parameters and persistent buffers (e.g. running averages) are
        included. Keys are corresponding parameter and buffer names.
        Parameters and buffers set to ``None`` are not included.

        Returns:
            dict:
                a dictionary containing a whole state of the module

        Example::

            >>> module.state_dict().keys()
            ['bias', 'weight']

        """
        if destination is None:
            destination = OrderedDict()
            destination._metadata = OrderedDict()
        destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
        self._save_to_state_dict(destination, prefix, keep_vars)
        for name, module in self._modules.items():
            if module is not None:
                module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
        for hook in self._state_dict_hooks.values():
            hook_result = hook(self, destination, prefix, local_metadata)
            if hook_result is not None:
                destination = hook_result
        return destination