commit change

mosaicml · Jun 3, 2024 · 5e3dadf · 5e3dadf
1 parent c3f74e5
commit 5e3dadf
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/data/test_icl_datasets.py b/tests/data/test_icl_datasets.py
@@ -37,15 +37,15 @@ def run_test(
     for i, e in enumerate(evaluators):
         batch = next(e.dataloader.dataloader.__iter__())
         # Check that the dataloader is the correct length for the first task.
-        if i == 0 and hasattr(e.dataloader.dataloader, '__len__'):
+        if i == 0:
             if eval_drop_last:
                 assert len(
                     e.dataloader.dataloader,
-                ) == EXPECTED_FIRST_DATALOADER_LEN - 1
+                ) == EXPECTED_FIRST_DATALOADER_LEN - 1  # pyright: ignore
             else:
                 assert len(
                     e.dataloader.dataloader,
-                ) == EXPECTED_FIRST_DATALOADER_LEN
+                ) == EXPECTED_FIRST_DATALOADER_LEN  # pyright: ignore
 
         inputs = batch['input_ids'][0]
         if 'continuation_indices' in batch: