Add missing **kwargs in on_epoch_begin of knowledge distill policy

distiller/knowledge_distillation.py

@@ -126,7 +126,7 @@ class KnowledgeDistillationPolicy(ScheduledTrainingPolicy):
 
     # Since the "forward" function isn't a policy callback, we use the epoch callbacks to toggle the
     # activation of distillation according the schedule defined by the user
-    def on_epoch_begin(self, model, zeros_mask_dict, meta):
+    def on_epoch_begin(self, model, zeros_mask_dict, meta, **kwargs):
         self.active = True
 
     def on_epoch_end(self, model, zeros_mask_dict, meta):