Fix KD integration test

Summary: D17271763 breaks integration test. This is a fix by adding the optional argument in KD task initializer. Reviewed By: akinh Differential Revision: D17431539 fbshipit-source-id: e1de2a6e2e5d51605922584616422aeb5d7b4944
pytorch · Sep 18, 2019 · c0c8867 · c0c8867
1 parent 46be5f9
commit c0c8867
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pytorch_translate/tasks/knowledge_distillation_task.py b/pytorch_translate/tasks/knowledge_distillation_task.py
@@ -19,12 +19,15 @@
 
 @register_task(constants.KNOWLEDGE_DISTILLATION_TASK)
 class PytorchKnowledgeDistillationTask(PytorchTranslateTask):
-    def __init__(self, args, src_dict, tgt_dict, char_source_dict=None):
+    def __init__(
+        self, args, src_dict, tgt_dict, char_source_dict=None, char_target_dict=None
+    ):
         super().__init__(
             args,
             src_dict=src_dict,
             tgt_dict=tgt_dict,
             char_source_dict=char_source_dict,
+            char_target_dict=char_target_dict,
         )
         self.top_k_probs_binary_file = args.top_k_probs_binary_file
         self.top_k_teacher_tokens = args.top_k_teacher_tokens