add model_ema to param_list

ain-soph · Nov 28, 2021 · 55719a9 · 55719a9
1 parent ebb8d58
commit 55719a9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/trojanzoo/trainer.py b/trojanzoo/trainer.py
@@ -22,7 +22,9 @@
 class Trainer:
     name = 'trainer'
     param_list = ['optim_args', 'train_args', 'writer_args',
-                  'optimizer', 'lr_scheduler', 'pre_conditioner', 'writer']
+                  'optimizer', 'lr_scheduler',
+                  'pre_conditioner', 'model_ema',
+                  'writer']
 
     @classmethod
     def add_argument(cls, group: argparse._ArgumentGroup):