How to combine multiple lightning module and save hyperparameters
Good day, I'm currently working on two models which train on the same data. I'd like to integrate the two pre-trained models into one and use it for transfer learning. The combination is written as such (you can copy paste to run it).
import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

class MyModelA(pl.LightningModule):
    def __init__(self, hidden_dim = 10):
        super(MyModelA, self).__init__()
        self.fc1 = torch.nn.Linear(hidden_dim, 2)
        self.save_hyperparameters()

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr = 1e-3)
        return optimizer
        
    def forward(self, x):
        x = self.fc1(x)
        return x

    def training_step(self, batch, batch_idx):
        x,y = batch
        return F.mse_loss(self.forward(x), y)
    
class MyModelB(pl.LightningModule):
    def __init__(self, hidden_dim = 10):
        super(MyModelB, self).__init__()
        self.fc1 = torch.nn.Linear(hidden_dim, 2)
        self.save_hyperparameters()
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr = 1e-3)
        return optimizer
        
    def forward(self, x):
        x = self.fc1(x)
        return x

    def training_step(self, batch, batch_idx):
        x,y = batch
        return F.mse_loss(self.forward(x), y)

class MyEnsemble(pl.LightningModule):
    def __init__(self, modelA, modelB):
        super(MyEnsemble, self).__init__()
        self.modelA = modelA
        self.modelB = modelB
        self.modelA.freeze()
        self.modelB.freeze()
        self.classifier = torch.nn.Linear(4, 2)

        #self.save_hyperparameters() # Uncomment to show error

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr = 1e-3)
        return optimizer
        
    def forward(self, x):
        x1 = self.modelA(x)
        x2 = self.modelB(x)
        x = torch.cat((x1, x2), dim=1)
        x = self.classifier(x)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        return F.mse_loss(self.forward(x), y)

dl = DataLoader(TensorDataset(torch.randn(1000, 10), 
                              torch.randn(1000, 2)), 
                batch_size = 10)

modelA = MyModelA()
modelB = MyModelB()

# pretrained modelA and modelB
trainerA = pl.Trainer(gpus = 0, max_epochs = 5, progress_bar_refresh_rate = 50)
trainerA.fit(modelA, dl)
trainerB = pl.Trainer(gpus = 0, max_epochs = 5, progress_bar_refresh_rate = 50)
trainerB.fit(modelB, dl)

# modelA and modelB contains pretrained weights
model = MyEnsemble(modelA, modelB)

trainer = pl.Trainer(gpus = 0, max_epochs = 5, progress_bar_refresh_rate = 50)
trainer.fit(model, dl)
At first the code worked fine. However, I would like to save the hyperparameters of the ensemble module, but adding self.save_hyperparameters() at the end of the ensemble module __init__ return this error.
ValueError: dictionary update sequence element #0 has length 1; 2 is required

Hence my question is how can I combine two or more lightning modules in a single module and save its hyperparameters? Or is there any alternative way to do so?
Thanks in advance!
EDIT: Code updated to show both modelA and modelB are pretrained.
I have finally came out with the final solution which can be obtained here.
Thank you for anyone who read and participate in this discussion.