How to pass losses to optimizer_step?

[AnkushMalaker]

I'm trying to use this multi-term Adam optimizer which requires me pass an array of losses to optimizer.step() as shown:

loss_array = unbalanced_loss(output, target, loss_weights)
ranks = [1]*10
optimizer.step(loss_array, ranks, None)

loss.backwards() is also called within the optimizer code. Here's what my attempt looks like:

1. Override backward() in model code and simply pass to disable loss.backward() by lightning.
2. Override the optimizer_step() hook.

In step 2, I need tp pass the losses to the optimizer_step() function. How can I do that, or a better way to achieve the same?
Thank you.

[rohitgr7]

i'd rather suggest using manual optimization to simplify your case since within Lightning, the optimizer_closure handles the training step and backward step so during actual optimizer_step, you might not have the losses available.

reF: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#pytorch_lightning.core.lightning.LightningModule.optimizer_step.params.optimizer_closure

ref: https://pytorch-lightning.readthedocs.io/en/stable/common/optimization.html#manual-optimization