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:

predictions = model(images)

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