Fix load_dataset signature (#1281)

Summary: Fix for #1240 Tested with MaskedLMTask. Pull Request resolved: #1281 Differential Revision: D18051472 fbshipit-source-id: 0aeff60c71489655f5e621349f780ba9cd8c027a
facebookresearch · Oct 22, 2019 · 34e6a5e · 34e6a5e
1 parent 66d24dc
commit 34e6a5e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/fairseq/tasks/masked_lm.py b/fairseq/tasks/masked_lm.py
@@ -71,7 +71,7 @@ def setup_task(cls, args, **kwargs):
         print('| dictionary: {} types'.format(len(dictionary)))
         return cls(args, dictionary)
 
-    def load_dataset(self, split, epoch=0, combine=False):
+    def load_dataset(self, split, epoch=0, combine=False, **kwargs):
         """Load a given dataset split.
 
         Args:

diff --git a/fairseq/tasks/multilingual_masked_lm.py b/fairseq/tasks/multilingual_masked_lm.py
@@ -110,7 +110,7 @@ def _get_sample_prob(self, dataset_lens):
         smoothed_prob = smoothed_prob / smoothed_prob.sum()
         return smoothed_prob
 
-    def load_dataset(self, split, epoch=0, combine=False):
+    def load_dataset(self, split, epoch=0, combine=False, **kwargs):
         """Load a given dataset split.
 
         Args: