diff --git a/finetune/trainer.py b/finetune/trainer.py
index 7da95ed..51f7f76 100644
--- a/finetune/trainer.py
+++ b/finetune/trainer.py
@@ -150,7 +150,7 @@ def prediction_step(
                 else:
                     loss = None
                     with self.compute_loss_context_manager():
-                        outputs = model(**inputs)
+                        outputs = model(data=inputs)
                     if isinstance(outputs, dict):
                         logits = tuple(
                             v for k, v in outputs.items() if k not in ignore_keys