From 2b9635c1e80448cb28339183e21de79571c0e4a8 Mon Sep 17 00:00:00 2001 From: Yan haixu <40758050+hova88@users.noreply.github.com> Date: Fri, 22 Sep 2023 11:26:39 +0800 Subject: [PATCH] add last_epoch in CosineAnnealingWarmupLR --- colossalai/nn/lr_scheduler/cosine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/colossalai/nn/lr_scheduler/cosine.py b/colossalai/nn/lr_scheduler/cosine.py index a896d3acba6c..f563825de0d5 100644 --- a/colossalai/nn/lr_scheduler/cosine.py +++ b/colossalai/nn/lr_scheduler/cosine.py @@ -62,7 +62,7 @@ def __init__(self, optimizer, total_steps: int, warmup_steps: int = 0, eta_min: base_scheduler = _CosineAnnealingLR( optimizer, total_steps - warmup_steps, eta_min=eta_min, last_epoch=last_epoch ) - super().__init__(optimizer, warmup_steps, base_scheduler) + super().__init__(optimizer, warmup_steps, base_scheduler, last_epoch=last_epoch) class FlatAnnealingLR(DelayerScheduler):