Internal change

PiperOrigin-RevId: 321720155

Internal change
PiperOrigin-RevId: 321720155
bf069622 · A. Unique TensorFlower · 0069c932 · bf069622
Commit bf069622 authored Jul 16, 2020 by A. Unique TensorFlower
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 2 deletions

official/nlp/tasks/masked_lm.py official/nlp/tasks/masked_lm.py +3 -2

No files found.
--- a/official/nlp/tasks/masked_lm.py
+++ b/official/nlp/tasks/masked_lm.py
@@ -40,8 +40,9 @@ class MaskedLMConfig(cfg.TaskConfig):
 class MaskedLMTask(base_task.Task):
  """Mock task object for testing."""
-  def build_model(self):
+  def build_model(self, params=None):
-    return bert.instantiate_pretrainer_from_cfg(self.task_config.model)
+    params = params or self.task_config.model
+    return bert.instantiate_pretrainer_from_cfg(params)
  def build_losses(self,
                   labels,