[misc] add last_epoch in CosineAnnealingWarmupLR (#4778)

hpcaitech · Sep 26, 2023 · a227063 · a227063
1 parent b6cf0ac
commit a227063
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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):