Best practice of adaptively updating parameters

[chris-tkinter]

Hi all!

I am trying to implement a model that is optimized using the following loss

loss = loss_A  + weight * loss_B

and would like to create a “weight warmup”, that is, the parameter weight starts at 0, and goes up linearly every epoch.

I know that the LightningModule has a self.current_epoch which is a counter for the epoch, so that I can manually increase the weight for every epoch. But I am just wondering is it the best practice?

Thanks!