debugging / 2

Helsinki-NLP · Sep 27, 2023 · ca5628d · ca5628d
1 parent 4476b26
commit ca5628d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/mammoth/inputters/dataloader.py b/mammoth/inputters/dataloader.py
@@ -191,8 +191,9 @@ def __iter__(self):
                 while self.bucket_is_empty(*current_bucket_idx):
                     current_bucket_idx = next(next_indices)
                 # retrieve and process the example
-                example = self._buckets[current_bucket_idx].pop()
-                self._lens[current_bucket_idx] -= 1
+                s_bucket, t_bucket = current_bucket_idx
+                example = self._buckets[s_bucket][t_bucket].pop()
+                self._lens[s_bucket][t_bucket] -= 1
                 accum.append(example)
                 numel = self.numel_fn(example)
                 cur_batch_size += numel