Update ofa_task.py

OFA-Sys · Jun 23, 2022 · 630e193 · 630e193
1 parent a742ad2
commit 630e193
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tasks/ofa_task.py b/tasks/ofa_task.py
@@ -132,7 +132,7 @@ def get_batch_iterator(
         total_row_count = dataset.dataset.get_total_row_count()
         num_batches = math.ceil(math.ceil(total_row_count / num_shards) / max_sentences)
         if len(batch_sampler) < num_batches:
-            batch_sampler.append([])
+            batch_sampler.append([1])
 
         # return a reusable, sharded iterator
         epoch_iter = iterators.EpochBatchIterator(