From 1444b1011a0dc35ec43bf32661dfd335ee0b1f29 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 30 Jun 2025 07:03:36 -0700 Subject: [PATCH] Update vision.py --- unsloth/models/vision.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index f0b0bdff1e..788df4c74e 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -435,8 +435,20 @@ class FastBaseModel: # Edit data-types if custom_datatype is not None: - for name, module in model.named_modules(): + for jj, (name, module) in enumerate(model.named_modules()): exec(custom_datatype) + if jj % 10 == 0: + gc.collect() + if DEVICE_TYPE == "cuda": torch.cuda.empty_cache() + elif DEVICE_TYPE == "xpu": torch.xpu.empty_cache() + pass + pass + # Clear deleted GPU items + for _ in range(3): + gc.collect() + if DEVICE_TYPE == "cuda": torch.cuda.empty_cache() + elif DEVICE_TYPE == "xpu": torch.xpu.empty_cache() + pass pass # Counteract saved tokenizers