github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-03 10:06:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] fix lr scheduler bug in torch 2.0 (#4864)

Browse Source
This commit is contained in:
Baizhou Zhang
2023-10-12 14:04:24 +08:00
committed by GitHub
parent 83b52c56cd
commit 39f2582e98
3 changed files with 17 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/nn/lr_scheduler/delayed.py
View File

@@ -1,4 +1,10 @@
from torch.optim.lr_scheduler import _LRScheduler
import torch
from packaging.version import Version
if Version(torch.__version__) >= Version("2.0.0"):
from torch.optim.lr_scheduler import LRScheduler as _LRScheduler
else:
from torch.optim.lr_scheduler import _LRScheduler
class _enable_get_lr_call: