How to use more than one optimizer at each step (jointly train multiple modules within one model)?
❓ Questions and Help
What is your question?
I have a model which consists of two blocks, let's call them first_module and second_module.
Code (simplified)
Training Step
def training_step(self, batch, batch_idx, optimizer_idx):
    out = self.first_module(batch)
    out = self.second_module(out)

    loss = criterion(out, batch['target'])

    metrics = {'train_loss': loss}
    output = {'loss': loss,
                  'log': metrics,
                  'progress_bar': metrics}

    return output
Optimizers
def configure_optimizers(self):
    train_params = self.train_params

    optimizer_first_module = torch.optim.Adam(
        self.first_module.parameters(),
        lr=train_params['lr_first_module'], betas=(0.5, 0.999))

    optimizer_second_module = torch.optim.Adam(
        self.second_module.parameters(),
        lr=train_params['lr_second_module'], betas=(0.5, 0.999))

    return [optimizer_first_module, optimizer_second_module]
Question
How to do optimizer_first_module.step() and optimizer_second_module.step() at each batch and ignore batch_idx?
It may be seen that optimizer_step always passes only one optimizer per step
    def optimizer_step(
        ...
        optimizer: Optimizer,
        optimizer_idx: int,
        ...
    ) -> None:
Possible solution (?)

Blending two optimizers into one (hacky way, not sure if pl would see the result as a correct optimizer class)
Modify training loop in PL (this option is even worse)