LightningModule __init__ vs. setup method

but also, you can use datamodules to decouple the data stuff from your model whenever you want. then the use of setup becomes more clear there.

DataModule:

  • setup
  • prepare data
  • train loader
  • val loader
  • test loader

LightningModule:

  • init
  • train step
  • val step (optional)
  • test step (optional)
  • optimizers
  • (optional forward)