fix test expectations (add line index support)

Helsinki-NLP · May 20, 2024 · ec347bc · ec347bc
1 parent 441f13c
commit ec347bc
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mammoth/tests/test_look_ahead_bucketing.py b/mammoth/tests/test_look_ahead_bucketing.py
@@ -1,5 +1,5 @@
 import pytest
-from itertools import product
+from itertools import product, count
 
 from mammoth.inputters.dataloader import build_dataloader
 
@@ -22,7 +22,7 @@ def __getitem__(self, idx):
     def __iter__(self):
         return iter(self.items)
 
-    def collate_fn(self, items):
+    def collate_fn(self, items, line_idx):
         return items