Fix remainder handling in chunking test

Signed-off-by: Alex-Brooks <[email protected]>
gkumbhat · Sep 29, 2023 · ccd9c13 · ccd9c13
1 parent 511e561
commit ccd9c13
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/resources/test_pretrained_model.py b/tests/resources/test_pretrained_model.py
@@ -166,7 +166,8 @@ def test_causal_lm_tok_output_correctness(models_cache_dir, chunk_size, drop_rem
     has_remainder = False
     if len(concat_tok) > chunk_size:
         num_expected_chunks = len(concat_tok) // chunk_size
-        if num_expected_chunks * chunk_size != len(concat_tok):
+        # Should only care about the remainder if we are not dropping it
+        if num_expected_chunks * chunk_size != len(concat_tok) and not drop_remainder:
             has_remainder = True
     else:
         num_expected_chunks = 1