Py typing fix

Googler · LIT team · commit d70e3d3c6467 · 2024-03-13T17:22:47.000-07:00
PiperOrigin-RevId: 615594853
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(
