"Forward vs training_step" section in style guide

Hi,

I am new to PyTorch Lightning and I cannot understand following parts in “Forward vs training_step” section in style guide:

We recommend using forward for inference/predictions and keeping training_step independent

def forward(...):
    embeddings = self.encoder(x)

def training_step(...):
    x, y = ...
    z = self.encoder(x)
    pred = self.decoder(z)
    ...

Isn’t single source of truth a virtue in programing? Why don’t you recommend reusing forward() in the training_step()?

you can use either. It’s just that sometimes inference can be bit different than training so forward should be defined according to what you want to infer from the model after training.
You can use this too:

def forward(...):
    embeddings = self.encoder(x)
    return embeddings

def training_step(...):
    x, y = ...
    z = self(x)
    pred = self.decoder(z)
    ...
1 Like