Generate different Outputs from model during training & during eval

Hello .

How can I make my model output different results while model is in train mode and while model is in eval.
For instance let’s say during train I would want my model to output the loss while during eval I would want my model to directly output the results.

I was looking for the functionality that present in the fasterrcnn_resnet50_fpn model wherein
The behavior of the model changes depending if it is in training or evaluation mode.

To achieve such behaviour, use the (special) self.training flag that is True when the model (the model has to subclass the torch.nn.Module) is in the train mode (model.train()) and False in the model’s eval mode (model.eval()).

And this is a skeleton of the forward pass of the model that outputs different results in the train/eval mode.

def forward(self, x):
    ...
    if self.training:
        ...
        return train_output
    else:
        ...
        return eval_output

@mariosasko Thanks so much for the solution.