def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False):
        """Save dict with state dicts of each of the models."""
        state_dict_ = {}
        if self.use_query_model:
            state_dict_[self._query_key] \
                = self.query_model.state_dict_for_save_checkpoint(
                    prefix=prefix, keep_vars=keep_vars)

        if self.use_block_model:
            state_dict_[self._block_key] \
                = self.block_model.state_dict_for_save_checkpoint(
                    prefix=prefix, keep_vars=keep_vars)

        return state_dict_
Save dict with state dicts of each of the models.
training/DeepSpeed-Domino/megatron/model/realm_model.py