bugfix

SeanLee97 · Sep 30, 2024 · 3b205f6 · 3b205f6
1 parent 5663a0f
commit 3b205f6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/angle_emb/angle.py b/angle_emb/angle.py
@@ -859,7 +859,8 @@ def compute_loss(self, model, inputs, return_outputs=False):
 
         return (loss, outputs) if return_outputs else loss
 
-    def prediction_step(self, model, inputs, **kwargs):
+    def prediction_step(self, model, inputs, *args, **kwargs):
+        kwargs.pop('prediction_loss_only', None)
         inputs['labels'] = None  # avoid evaluation error
         return super().prediction_step(model, inputs, prediction_loss_only=True, **kwargs)