github
/
minimind-v
mirror of https://github.com/jingyaogong/minimind-v.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[feat] del cuda memory

pull/107/head
jingyaogong 2025-11-27 19:37:39 +08:00
parent e613613a13
commit a157b66ace
2 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
trainer/train_pretrain_vlm.py
View File

@ -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")

3
trainer/train_sft_vlm.py
View File

@ -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")