From 0477c5804e4d2a03480b81936c6ac1a98f391fdb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 22 Jul 2025 03:34:51 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index de53704e66..9b920fdac9 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -86,16 +86,13 @@ from triton import __version__ as triton_version HAS_XFORMERS = xformers is not None BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None -def clean_gpu_cache(): - if DEVICE_TYPE == "xpu": - torch.xpu.empty_cache() - else: - torch.cuda.empty_cache() - if DEVICE_TYPE == "xpu": + clean_gpu_cache = torch.xpu.empty_cache get_current_device = torch.xpu.current_device else: + clean_gpu_cache = torch.cuda.empty_cache get_current_device = torch.cuda.current_device +pass def original_apply_qkv(self, X): Q = self.q_proj(X)