How to save finnetunned weights when using two transformers BERT models inside Pytorch Lightning

[math-sasso]

In my case I want to save weight for question_encoder AND for doc_encoder

class TwoTower(pl.LightningModule):
    """Two tower model for document retrieval using query and document."""
    def __init__(self,
                 hparams):
        super(TwoTower, self).__init__()
        self.hparams = hparams
        pl.utilities.seed.seed_everything(seed = self.hparams.seed)

    def setup(self, step):
        self.training = False

        # Models
        self.question_encoder = Encoder(dim=self.hparams.dim, pretrained_model=self.hparams.model_name)
        self.doc_encoder = Encoder(dim=self.hparams.dim, pretrained_model=self.hparams.model_name)

        # Loss
        self.faiss_module = FaissModule()