From 597953bef0245d551b4d23c715d3ae2b2a7dd549 Mon Sep 17 00:00:00 2001 From: VED <146507396+ved1beta@users.noreply.github.com> Date: Thu, 7 Aug 2025 18:55:58 +0530 Subject: [PATCH] clear cache before clean up (#3031) [skip ci] * clear chahe before save_model * chore: lint --------- Co-authored-by: Wing Lian --- src/axolotl/train.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/axolotl/train.py b/src/axolotl/train.py index 41f184abc..a693236d3 100644 --- a/src/axolotl/train.py +++ b/src/axolotl/train.py @@ -566,6 +566,10 @@ def train( resume_from_checkpoint = determine_resume_checkpoint(cfg) 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_trained_model(cfg, trainer, model, safe_serialization) create_model_card(cfg, trainer)