Log EMA decay

1 files changed, 2 insertions, 1 deletions

diff --git a/train_ti.py b/train_ti.py
index dc36e42..2f13128 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -270,7 +270,6 @@ def parse_args():
     parser.add_argument(
         "--use_ema",
         action="store_true",
-        default=True,
         help="Whether to use EMA model."
     )
     parser.add_argument(
@@ -1004,6 +1003,8 @@ def main():
                     "train/cur_acc": acc.item(),
                     "lr": lr_scheduler.get_last_lr()[0],
                 }
+                if args.use_ema:
+                    logs["ema_decay"] = ema_embeddings.decay
 
                 accelerator.log(logs, step=global_step)