custom prediction step

SeanLee97 · Sep 30, 2024 · 5663a0f · 5663a0f
1 parent a07f045
commit 5663a0f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/angle_emb/angle.py b/angle_emb/angle.py
@@ -859,6 +859,10 @@ def compute_loss(self, model, inputs, return_outputs=False):
 
         return (loss, outputs) if return_outputs else loss
 
+    def prediction_step(self, model, inputs, **kwargs):
+        inputs['labels'] = None  # avoid evaluation error
+        return super().prediction_step(model, inputs, prediction_loss_only=True, **kwargs)
+
 
 class AngleESETrainer(AngleTrainer):
     """