fix test_ma

marrlab · Dec 5, 2024 · 3cef1a2 · 3cef1a2
1 parent b296229
commit 3cef1a2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/domainlab/algos/trainers/train_ema.py b/domainlab/algos/trainers/train_ema.py
@@ -63,4 +63,4 @@ def after_epoch(self, epoch, flag_info=None):
         new_dict_para = self.move_average(dict_para, epoch)
         # without deepcopy, this seems to work
         torch_model.load_state_dict(new_dict_para)
-        super().after_epoch(epoch)
+        super().after_epoch(epoch, flag_info)