Py typing fix

PiperOrigin-RevId: 615594853
PAIR-code · Mar 14, 2024 · d70e3d3 · d70e3d3
1 parent 45887d3
commit d70e3d3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lit_nlp/components/cf_utils.py b/lit_nlp/components/cf_utils.py
@@ -40,7 +40,8 @@ def update_prediction(
     label_names = pred_spec.vocab
     pred_class = np.argmax(prediction)
     example_label = label_names[pred_class]
-    example[label_key] = example_label
+    if label_key is not None:
+      example[label_key] = example_label
 
 
 def is_prediction_flip(