Log DAdam/DAdan d

1 files changed, 14 insertions, 0 deletions

diff --git a/training/functional.py b/training/functional.py
index 43ee356..77f056e 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -424,6 +424,15 @@ def train_loop(
     on_sample = callbacks.on_sample
     on_checkpoint = callbacks.on_checkpoint
 
+    isDadaptation = False
+
+    try:
+        import dadaptation
+
+        isDadaptation = isinstance(optimizer, (dadaptation.DAdaptAdam, dadaptation.DAdaptAdan))
+    except ImportError:
+        pass
+
     try:
         for epoch in range(num_epochs):
             if accelerator.is_main_process:
@@ -461,6 +470,11 @@ def train_loop(
                         "train/cur_acc": acc.item(),
                         "lr": lr_scheduler.get_last_lr()[0],
                     }
+                    if isDadaptation:
+                        logs["lr/d*lr"] = (
+                            optimizer.param_groups[0]["d"] *
+                            optimizer.param_groups[0]["lr"]
+                        )
                     logs.update(on_log())
 
                     local_progress_bar.set_postfix(**logs)