

TITLE
prediction_step() is not using compute_loss()

COMMENTS
2

REACTIONS
+1: 1
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
Hi @sgugger , I think there is an issue with `prediction_step` in `trainer.py`. The problem arises when implementing a custom loss function that requires reshaping input labels. In `training_step`, for loss calculation, it calls `compute_loss()` which is totally fine but in the `prediction_step` it calculates loss without `compute_loss()` function. This inconsistency causes some issues. Would you think it would be better to call `compute_loss()` in both cases in order to avoid this problem? 

Update: As I see the code, the loss is tightly coupled with `LabelSmoother` which makes it hard to do them in a single function. If you have any suggestion, it makes me happy to contribute to Huggingface ;)

[prediction_step()](https://huggingface.co/transformers/_modules/transformers/trainer.html#Trainer.prediction_step)

[training_step()](https://huggingface.co/transformers/_modules/transformers/trainer.html#Trainer.training_step)