Can pylightning handle multiple separate losses?

[Mi-Przystupa]

I was curious to know if / how to have pylightning make multiple backward passes through losses that are to be done separate?

I'm trying to set-up soft actor critic as a PyLightning Module & the policy update and q_value update shouldn't be combined.

So in my train step:

def train_step(self, batch):
   # Rl update logics
   return loss1, loss2

[Mi-Przystupa]

Just skimming the documents:
https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html?highlight=training_step#training-step

I guess defining multiple optimizer's is exactly what I'm looking for :

# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx, optimizer_idx):
    if optimizer_idx == 0:
        # do training_step with encoder
        ...
    if optimizer_idx == 1:
        # do training_step with decoder
        ...