How to change between train and evaluate mode in forward method?

I’m trying to build a next frame generation model which behaves in the following manner"

  • During training, the model expects both the input tensors, as well as targets and will return a dict containing the classification and regression losses:
loss_dict = model(images, targets)
  • During evaluation, the model requires only the input tensors, and returns the post-processed
predictions = model(images)

How can I toggle between these two modes in def forward(self, images): of a PyTorch Lightning system?

Each nn.Module uses the internal flag, which is changed by calling model.train() and model.eval(), and can be used as a condition to use different behaviors inside the module.

Thanks I understand.