From 94643fac8a27345f695500085d78cc8fa01f5fa9 Mon Sep 17 00:00:00 2001 From: Leo Jiang <74156916+leisuzz@users.noreply.github.com> Date: Fri, 25 Oct 2024 08:35:19 -0600 Subject: [PATCH] [bugfix] bugfix for npu free memory (#9640) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Improve NPU performance * Improve NPU performance * Improve NPU performance * Improve NPU performance * [bugfix] bugfix for npu free memory * [bugfix] bugfix for npu free memory * [bugfix] bugfix for npu free memory --------- Co-authored-by: 蒋硕 Co-authored-by: Sayak Paul --- src/diffusers/training_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/training_utils.py b/src/diffusers/training_utils.py index 9c898ad141..0e0d0ce5b5 100644 --- a/src/diffusers/training_utils.py +++ b/src/diffusers/training_utils.py @@ -284,7 +284,7 @@ def free_memory(): elif torch.backends.mps.is_available(): torch.mps.empty_cache() elif is_torch_npu_available(): - torch_npu.empty_cache() + torch_npu.npu.empty_cache() # Adapted from torch-ema https://github.com/fadel/pytorch_ema/blob/master/torch_ema/ema.py#L14