Fix NVIDIA#349

Signed-off-by: Atsunori Fujita <[email protected]>
AtsunoriFujita · Oct 18, 2024 · dd9954b · dd9954b
1 parent d8ef9fb
commit dd9954b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/examples/nlp/data/steerlm/attribute_annotate.py b/examples/nlp/data/steerlm/attribute_annotate.py
@@ -71,7 +71,8 @@ def get_reward(
         reward_out = output_dict["rewards"].flatten().tolist()
 
         all_rewards.append(reward_out)
-        all_exceeded += output_dict["exceeded"].tolist()
+        if output_dict.get('exceeded', None) is not None:
+            all_exceeded += output_dict["exceeded"].tolist()
 
     return all_rewards, all_exceeded