fix the memory leak of offloading

AlibabaPAI · Dec 12, 2024 · 3eb1a89 · 3eb1a89
1 parent 88cee13
commit 3eb1a89
Showing 1 changed file with 5 additions and 1 deletion.
diff --git 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),