pop labels

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