1 files changed, 13 insertions, 5 deletions

diff --git a/training/functional.py b/training/functional.py
index ac43847..7104a88 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -484,12 +484,16 @@ def train_loop(
                     avg_loss.update(loss.detach_(), bsz)
                     avg_acc.update(acc.detach_(), bsz)
 
+                    lr = lr_scheduler.get_last_lr()[0]
+                    if torch.is_tensor(lr):
+                        lr = lr.item()
+
                     logs = {
                         "train/loss": avg_loss.avg.item(),
                         "train/acc": avg_acc.avg.item(),
                         "train/cur_loss": loss.item(),
                         "train/cur_acc": acc.item(),
-                        "lr": lr_scheduler.get_last_lr()[0],
+                        "lr": lr,
                     }
                     if isDadaptation:
                         logs["lr/d*lr"] = optimizer.param_groups[0]["d"] * optimizer.param_groups[0]["lr"]
@@ -498,13 +502,13 @@ def train_loop(
                     local_progress_bar.set_postfix(**logs)
 
                     if ((step + 1) % gradient_accumulation_steps == 0) or ((step + 1) == len(train_dataloader)):
-                        before_optimize_result = on_before_optimize(lr_scheduler.get_last_lr()[0], epoch)
+                        before_optimize_result = on_before_optimize(lr, epoch)
 
                         optimizer.step()
                         lr_scheduler.step()
                         optimizer.zero_grad(set_to_none=True)
 
-                        on_after_optimize(before_optimize_result, lr_scheduler.get_last_lr()[0])
+                        on_after_optimize(before_optimize_result, lr)
 
                         local_progress_bar.update(1)
                         global_progress_bar.update(1)
@@ -518,9 +522,13 @@ def train_loop(
 
             accelerator.wait_for_everyone()
 
-            lrs.append(lr_scheduler.get_last_lr()[0])
+            lr = lr_scheduler.get_last_lr()[0]
+            if torch.is_tensor(lr):
+                lr = lr.item
+
+            lrs.append(lr)
 
-            on_after_epoch(lr_scheduler.get_last_lr()[0])
+            on_after_epoch(lr)
 
             if val_dataloader is not None:
                 model.eval()