Fix missing autocast() in Trainer.prediction_step() (#14075)

Co-authored-by: jonas <jonas@hpcnt.com>

src/transformers/trainer.py

@@ -2486,6 +2486,10 @@ class Trainer:
                     logits = smp_nested_concat(logits_mb)
             else:
                 if has_labels:
+                    if self.use_amp:
+                        with autocast():
+                            loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
+                    else:
                         loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
                     loss = loss.mean().detach()
                     if isinstance(outputs, dict):