[hotfix] fix for loading after_scheduler state #5120
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
馃搶 Checklist before creating the PR
[doc/gemini/tensor/...]: A concise description
馃毃 Issue number
related issue #5119
馃摑 What does this PR do?
I have introduced a new method, load_state_dict, in the Scheduler class in colossalai/nn/lr_scheduler/delayed.py. This method enhances the scheduler's ability to reinstate its state, specifically focusing on the after_scheduler component. The key changes include:
Updating the Scheduler class with the load_state_dict method.
Ensuring the after_scheduler's state is appropriately loaded from the provided state dictionary if the after_scheduler_dict key is present.
馃挜 Checklist before requesting a review
猸愶笍 Do you enjoy contributing to Colossal-AI?
Tell us more if you don't enjoy contributing to Colossal-AI.