[lr-scheduler] fix load state dict and add test (#5369)

tests/test_optimizer/test_lr_scheduler.py

@@ -0,0 +1,20 @@
+import torch.nn as nn
+from torch.optim import Adam
+
+from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
+
+
+def test_lr_scheduler_save_load():
+    model = nn.Linear(10, 10)
+    optimizer = Adam(model.parameters(), lr=1e-3)
+    scheduler = CosineAnnealingWarmupLR(optimizer, total_steps=5, warmup_steps=2)
+    new_scheduler = CosineAnnealingWarmupLR(optimizer, total_steps=5, warmup_steps=2)
+    for _ in range(5):
+        scheduler.step()
+        state_dict = scheduler.state_dict()
+        new_scheduler.load_state_dict(state_dict)
+        assert state_dict == new_scheduler.state_dict()
+
+
+if __name__ == "__main__":
+    test_lr_scheduler_save_load()