You can have your model return anything you want, but returning a dictionary like we did for
ResnetModelForImageClassification, with the loss included when labels are passed, will make your model directly
usable inside the [Trainer] class.