LightningModule __init__ vs. setup method

If there is anything in particular that you want to change when you are running in multiple stages then setup can be used. Like if you need to set up or get the model ready before testing, that can be done with setup('test')