I would like to perform some operations on the output probability distributions of an AutoModelForSequenceClassification model so, I was wondering if it is possible to return the logits rather than predicted class labels from the transformers.Trainer.predict() method.