roll back the changes in mm_utils.py

NVlabs · May 8, 2024 · 2a8bb83 · 2a8bb83
1 parent f85297f
commit 2a8bb83
Showing 1 changed file with 1 addition and 12 deletions.
diff --git a/llava/mm_utils.py b/llava/mm_utils.py
@@ -269,18 +269,7 @@ def call_for_batch(
             keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids
         ]
         for keyword_id in self.keyword_ids:
-            if keyword_id.ndim == 3:
-                if (output_ids[0, -keyword_id.shape[0] :, None] == keyword_id).all():
-                    return True
-            elif keyword_id.ndim == 2:
-                if (output_ids[0, -keyword_id.shape[0] :] == keyword_id).all():
-                    return True
-            else:
-                raise ValueError(
-                    "Keyword tensor should have 2 or 3 dimensions, got {}".format(
-                        keyword_id.ndim
-                    )
-                )
+            if (output_ids[0, -keyword_id.shape[0] :] == keyword_id).all():
                 return True
         outputs = self.tokenizer.batch_decode(
             output_ids[:, -offset:], skip_special_tokens=True