From 3eb1a895998ea9505cc993ca5c5ea56a5eb27e53 Mon Sep 17 00:00:00 2001 From: Yitong Huang Date: Thu, 12 Dec 2024 11:48:58 +0800 Subject: [PATCH] fix the memory leak of offloading --- torchacc/utils/cpu_offload.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/torchacc/utils/cpu_offload.py b/torchacc/utils/cpu_offload.py index 9c45718..61d55a3 100644 --- a/torchacc/utils/cpu_offload.py +++ b/torchacc/utils/cpu_offload.py @@ -593,7 +593,11 @@ def apply(tensor): _applied_tensor = tensor return tensor - return apply_to_tensors(apply, outputs) + outputs = apply_to_tensors(apply, outputs) + # delete the global variables to avoid memory leak + _apply_tensor = None + _applied_tensor = None + return outputs return ( CpuOffloadHookWithOffloadHandler(offload_handler=cpu_offload_handler),