From ec347bca0200cf7650b39748c2e8d01b3eb2ddb7 Mon Sep 17 00:00:00 2001 From: Mickus Timothee Date: Tue, 26 Mar 2024 09:21:16 +0200 Subject: [PATCH] fix test expectations (add line index support) --- mammoth/tests/test_look_ahead_bucketing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mammoth/tests/test_look_ahead_bucketing.py b/mammoth/tests/test_look_ahead_bucketing.py index 6a97f79f..7d40029f 100644 --- 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