@@ -34,6 +34,7 @@ class KDEMATest(unittest.TestCase):
             "greater_metric_to_watch_is_better": True,
             "average_best_models": False,
             "ema": True,
+            "ema_params": {"decay_type": "constant", "decay": 0.999},
         }
 
     def test_teacher_ema_not_duplicated(self):

@@ -133,6 +133,8 @@ class KDTrainerTest(unittest.TestCase):
         train_params = self.kd_train_params.copy()
         train_params["max_epochs"] = 1
         train_params["ema"] = True
+        train_params["ema_params"] = {"decay_type": "constant", "decay": 0.999}
+
         kd_trainer.train(
             training_params=train_params,
             student=student,