fix stopping_criteria check in coca_model

mlfoundations · Apr 12, 2024 · b3b5fea · b3b5fea
1 parent 9eaf242
commit b3b5fea
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/open_clip/coca_model.py b/src/open_clip/coca_model.py
@@ -298,7 +298,7 @@ def generate(
 
                 cur_len += 1
 
-                if stopping_criteria(out, None):
+                if stopping_criteria(out, None).any():
                     break
 
             if num_dims == 1:
@@ -439,7 +439,7 @@ def _generate_beamsearch(
 
             # increase cur_len
             cur_len = cur_len + 1
-            if beam_scorer.is_done or stopping_criteria(input_ids, None):
+            if beam_scorer.is_done or stopping_criteria(input_ids, None).any():
                 break
 
         final_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None