tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

clear cache before clean up (#3031) [skip ci]

Browse Source 
* clear chahe before save_model

* chore: lint

---------

Co-authored-by: Wing Lian <wing@axolotl.ai>
This commit is contained in:
VED
2025-08-07 18:55:58 +05:30
committed by GitHub
parent 39fbd3b2b5
commit 597953bef0
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/axolotl/train.py
View File

@@ -566,6 +566,10 @@ def train(
resume_from_checkpoint = determine_resume_checkpoint(cfg) resume_from_checkpoint = determine_resume_checkpoint(cfg)
execute_training(cfg, trainer, resume_from_checkpoint) execute_training(cfg, trainer, resume_from_checkpoint)
# clear cache
if torch.cuda.is_available():
torch.cuda.empty_cache()
# Save the trained model and cleanup # Save the trained model and cleanup
save_trained_model(cfg, trainer, model, safe_serialization) save_trained_model(cfg, trainer, model, safe_serialization)
create_model_card(cfg, trainer) create_model_card(cfg, trainer)