diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py
index 32b155d..b4997b2 100644
--- a/trainer/train_pretrain_vlm.py
+++ b/trainer/train_pretrain_vlm.py
@@ -82,6 +82,9 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
                          epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
             model.train()
 
+        del X, Y, loss_mask, pixel_values, res, loss
+        torch.cuda.empty_cache()
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind-V Pretrain")
diff --git a/trainer/train_sft_vlm.py b/trainer/train_sft_vlm.py
index 253436e..180b4b6 100644
--- a/trainer/train_sft_vlm.py
+++ b/trainer/train_sft_vlm.py
@@ -82,6 +82,9 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
                          epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
             model.train()
 
+        del X, Y, loss_mask, pixel_values, res, loss
+        torch.cuda.empty_cache()
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind-V SFT")